forked from tinygrad/tinygrad
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 geohot
|
3b8dad62f2 |
@@ -41,6 +41,10 @@ inputs:
|
||||
description: "Install LLVM?"
|
||||
required: false
|
||||
default: 'false'
|
||||
tinydreno:
|
||||
description: "Install tinydreno"
|
||||
required: false
|
||||
default: 'false'
|
||||
qemu:
|
||||
description: "Install qemu"
|
||||
required: false
|
||||
@@ -273,6 +277,12 @@ 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'
|
||||
|
||||
@@ -94,6 +94,7 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -148,6 +149,7 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -200,6 +202,7 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -249,6 +252,7 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
|
||||
@@ -1,181 +0,0 @@
|
||||
name: Platform Tests
|
||||
env:
|
||||
# increment this when downloads substantially change to avoid the internet
|
||||
CACHE_VERSION: '19'
|
||||
CAPTURE_PROCESS_REPLAY: ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.title, '[pr]') && '1' || '0' }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
PYTHONPATH: ${{ github.workspace }}
|
||||
CHECK_OOB: 1
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- master
|
||||
pull_request:
|
||||
workflow_dispatch:
|
||||
|
||||
concurrency:
|
||||
group: platform-${{ github.event_name }}-${{ github.event_name == 'pull_request' && github.event.pull_request.number || github.run_id }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
|
||||
# ****** OSX Tests ******
|
||||
|
||||
unittestmacos:
|
||||
name: MacOS (unit)
|
||||
runs-on: macos-26
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: unittest-macos
|
||||
deps: testing_unit
|
||||
- name: Run unit tests
|
||||
run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
|
||||
- name: Test tensor core ops (fake)
|
||||
run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
|
||||
- name: Test tensor core ops (real)
|
||||
run: DEV=METAL DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
|
||||
- name: Test Beam Search
|
||||
run: DEV=METAL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
|
||||
- name: Test Device Specific
|
||||
run: DEV=METAL python3 -m pytest test/device/test_metal.py
|
||||
#- name: Fuzz Test linearizer
|
||||
# run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
unittestmacosmock:
|
||||
name: MacOS (unit, mock)
|
||||
runs-on: macos-26
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: unittest-macos-mock
|
||||
deps: testing_unit
|
||||
amd: 'true'
|
||||
ocelot: 'true'
|
||||
- name: Run NULL backend tests
|
||||
run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
|
||||
- name: Run pytest (amd)
|
||||
env:
|
||||
DEV: MOCKKFD+AMD
|
||||
FORWARD_ONLY: 1
|
||||
run: |
|
||||
python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
|
||||
- name: Run pytest (ptx)
|
||||
env:
|
||||
DEV: "MOCK+NV:PTX"
|
||||
FORWARD_ONLY: 1
|
||||
# TODO: failing due to library loading error
|
||||
CAPTURE_PROCESS_REPLAY: 0
|
||||
run: |
|
||||
python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py \
|
||||
test/testextra/test_hevc.py::TestHevc::test_hevc_decode_compile --durations=20
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
testmetal:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
group: [1, 2]
|
||||
name: MacOS (DEV=METAL) (${{ matrix.group }})
|
||||
runs-on: macos-26
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
DEV: METAL
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: macos-metal
|
||||
deps: testing_unit
|
||||
- name: Check Device.DEFAULT and print some source
|
||||
run: |
|
||||
python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'"
|
||||
DEBUG=4 python test/test_tiny.py TestTiny.test_plus
|
||||
- name: Run backend tests
|
||||
run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }}
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
testmacos:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
dev:
|
||||
- 'CPU:CLANG'
|
||||
- 'CPU:LLVM'
|
||||
- 'CPU:LVP'
|
||||
- 'WEBGPU'
|
||||
|
||||
name: MacOS (DEV=${{ matrix.dev }})
|
||||
runs-on: macos-26
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: macos-${{ matrix.dev }}
|
||||
deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
|
||||
llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
|
||||
webgpu: ${{ matrix.dev == 'WEBGPU' }}
|
||||
- name: Set env
|
||||
run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
|
||||
- name: Check Device.DEFAULT and print some source
|
||||
run: |
|
||||
python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
|
||||
DEBUG=4 python test/test_tiny.py TestTiny.test_plus
|
||||
- name: Run test_tiny
|
||||
run: python -m pytest -n=auto test/test_tiny.py --durations=20
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
# ****** Windows Tests ******
|
||||
|
||||
testwindows:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
dev:
|
||||
- 'CPU:CLANG'
|
||||
- 'CPU:LLVM'
|
||||
- 'CPU:X86'
|
||||
- 'WEBGPU'
|
||||
|
||||
name: Windows (DEV=${{ matrix.dev }})
|
||||
runs-on: windows-2025
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: windows-${{ matrix.dev }}-minimal
|
||||
deps: testing_unit
|
||||
pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
|
||||
- name: Set env
|
||||
shell: bash
|
||||
run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
|
||||
- name: Check Device.DEFAULT and print some source
|
||||
shell: bash
|
||||
run: |
|
||||
python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
|
||||
DEBUG=4 python test/test_tiny.py TestTiny.test_plus
|
||||
- name: Run test_tiny
|
||||
shell: bash
|
||||
run: python -m pytest -n=auto test/test_tiny.py --durations=20
|
||||
+229
-39
@@ -21,7 +21,7 @@ concurrency:
|
||||
jobs:
|
||||
docs:
|
||||
name: Docs
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
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' }}
|
||||
timeout-minutes: 10
|
||||
env:
|
||||
CHECK_OOB: 0
|
||||
@@ -61,7 +61,7 @@ jobs:
|
||||
|
||||
torchbackend:
|
||||
name: Torch Backend Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -88,7 +88,7 @@ jobs:
|
||||
|
||||
bepython:
|
||||
name: Python Backend
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -126,7 +126,7 @@ jobs:
|
||||
|
||||
linter:
|
||||
name: Linters
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 10
|
||||
|
||||
steps:
|
||||
@@ -157,7 +157,7 @@ jobs:
|
||||
|
||||
nulltest:
|
||||
name: Null Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
@@ -191,7 +191,7 @@ jobs:
|
||||
|
||||
unittest:
|
||||
name: Unit Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
@@ -228,7 +228,7 @@ jobs:
|
||||
matrix:
|
||||
group: [1, 2]
|
||||
name: SPEC=2 (${{ matrix.group }})
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -244,7 +244,7 @@ jobs:
|
||||
|
||||
fuzzing:
|
||||
name: Fuzzing
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -260,7 +260,7 @@ jobs:
|
||||
|
||||
testopenclimage:
|
||||
name: CL IMAGE Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -280,7 +280,7 @@ jobs:
|
||||
|
||||
testopenpilot:
|
||||
name: openpilot Compile Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -309,7 +309,7 @@ jobs:
|
||||
|
||||
testonnxcpu:
|
||||
name: ONNX (CPU) Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 20
|
||||
|
||||
steps:
|
||||
@@ -328,7 +328,7 @@ jobs:
|
||||
|
||||
testoptim:
|
||||
name: Optimization Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -360,7 +360,7 @@ jobs:
|
||||
|
||||
testllm:
|
||||
name: Test LLM
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
CHECK_OOB: 0
|
||||
@@ -387,7 +387,7 @@ jobs:
|
||||
|
||||
testmodels:
|
||||
name: Models
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -407,7 +407,7 @@ jobs:
|
||||
|
||||
testdsp:
|
||||
name: Linux (DSP)
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -435,7 +435,7 @@ jobs:
|
||||
- 'WEBGPU'
|
||||
|
||||
name: Linux (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' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -461,7 +461,7 @@ jobs:
|
||||
|
||||
testamdasm:
|
||||
name: AMD ASM IDE
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
DEV: MOCKKFD+AMD
|
||||
@@ -507,7 +507,7 @@ jobs:
|
||||
|
||||
hcq2:
|
||||
name: hcq2
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -521,15 +521,16 @@ 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 -m pytest -n=auto test/backend/test_multitensor.py
|
||||
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
|
||||
- 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: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
DEV: MOCKPCI+AMD
|
||||
@@ -565,7 +566,7 @@ jobs:
|
||||
arch: [gfx1100, gfx1201, gfx950]
|
||||
|
||||
name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
|
||||
@@ -588,7 +589,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 --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 test/opt/test_tensor_cores.py --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
|
||||
@@ -603,7 +604,7 @@ jobs:
|
||||
backend: [ptx, nv]
|
||||
|
||||
name: Linux (${{ matrix.backend }})
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
FORWARD_ONLY: 1
|
||||
@@ -631,23 +632,175 @@ jobs:
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
# ****** OSX Tests ******
|
||||
|
||||
unittestmacos:
|
||||
name: MacOS (unit)
|
||||
runs-on: macos-26
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: unittest-macos
|
||||
deps: testing_unit
|
||||
- name: Run unit tests
|
||||
run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
|
||||
- name: Test tensor core ops (fake)
|
||||
run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
|
||||
- name: Test tensor core ops (real)
|
||||
run: DEV=METAL DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
|
||||
- name: Test Beam Search
|
||||
run: DEV=METAL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
|
||||
- name: Test Device Specific
|
||||
run: DEV=METAL python3 -m pytest test/device/test_metal.py
|
||||
#- name: Fuzz Test linearizer
|
||||
# run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
unittestmacosmock:
|
||||
name: MacOS (unit, mock)
|
||||
runs-on: macos-26
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: unittest-macos-mock
|
||||
deps: testing_unit
|
||||
amd: 'true'
|
||||
ocelot: 'true'
|
||||
- name: Run NULL backend tests
|
||||
run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
|
||||
- name: Run pytest (amd)
|
||||
env:
|
||||
DEV: MOCKKFD+AMD
|
||||
FORWARD_ONLY: 1
|
||||
run: |
|
||||
python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
|
||||
- name: Run pytest (ptx)
|
||||
env:
|
||||
DEV: "MOCK+NV:PTX"
|
||||
FORWARD_ONLY: 1
|
||||
# TODO: failing due to library loading error
|
||||
CAPTURE_PROCESS_REPLAY: 0
|
||||
run: |
|
||||
python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
testmetal:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
group: [1, 2]
|
||||
name: MacOS (DEV=METAL) (${{ matrix.group }})
|
||||
runs-on: macos-26
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
DEV: METAL
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: macos-metal
|
||||
deps: testing_unit
|
||||
- name: Check Device.DEFAULT and print some source
|
||||
run: |
|
||||
python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'"
|
||||
DEBUG=4 python test/test_tiny.py TestTiny.test_plus
|
||||
- name: Run backend tests
|
||||
run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }}
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
testmacos:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
dev:
|
||||
- 'CPU:CLANG'
|
||||
- 'CPU:LLVM'
|
||||
- 'CPU:LVP'
|
||||
- 'WEBGPU'
|
||||
|
||||
name: MacOS (DEV=${{ matrix.dev }})
|
||||
runs-on: macos-26
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: macos-${{ matrix.dev }}
|
||||
deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
|
||||
llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
|
||||
webgpu: ${{ matrix.dev == 'WEBGPU' }}
|
||||
- name: Set env
|
||||
run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
|
||||
- name: Check Device.DEFAULT and print some source
|
||||
run: |
|
||||
python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
|
||||
DEBUG=4 python test/test_tiny.py TestTiny.test_plus
|
||||
- name: Run test_tiny
|
||||
run: python -m pytest -n=auto test/test_tiny.py --durations=20
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
# ****** Windows Tests ******
|
||||
|
||||
testwindows:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
dev:
|
||||
- 'CPU:CLANG'
|
||||
- 'CPU:LLVM'
|
||||
- 'CPU:X86'
|
||||
- 'WEBGPU'
|
||||
|
||||
name: Windows (DEV=${{ matrix.dev }})
|
||||
runs-on: windows-2025
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: windows-${{ matrix.dev }}-minimal
|
||||
deps: testing_unit
|
||||
pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
|
||||
- name: Set env
|
||||
shell: bash
|
||||
run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
|
||||
- name: Check Device.DEFAULT and print some source
|
||||
shell: bash
|
||||
run: |
|
||||
python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
|
||||
DEBUG=4 python test/test_tiny.py TestTiny.test_plus
|
||||
- name: Run test_tiny
|
||||
shell: bash
|
||||
run: python -m pytest -n=auto test/test_tiny.py --durations=20
|
||||
|
||||
# ****** Compile-only Tests ******
|
||||
|
||||
compiletests:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
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' }}
|
||||
backend: [ir3, nak]
|
||||
name: Compile-only (${{ matrix.backend }})
|
||||
runs-on: *linux
|
||||
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
|
||||
@@ -656,14 +809,51 @@ jobs:
|
||||
with:
|
||||
key: compile-${{ matrix.backend }}
|
||||
deps: "testing_unit mesa"
|
||||
qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
|
||||
- name: Test IMAGE
|
||||
- name: Set env
|
||||
shell: bash
|
||||
if: contains(matrix.dev, 'a630')
|
||||
run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
|
||||
run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $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)
|
||||
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
|
||||
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
|
||||
|
||||
@@ -1458,8 +1458,7 @@ def train_llama3():
|
||||
|
||||
# realize everything here
|
||||
if optim.master_params: Tensor.realize(*optim.master_params)
|
||||
loss_acc = Tensor.zeros(1, dtype=dtypes.float32, device=device)
|
||||
Tensor.realize(loss_acc, *optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
|
||||
Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
|
||||
|
||||
@TinyJit
|
||||
def minibatch(tokens:Tensor):
|
||||
@@ -1477,8 +1476,8 @@ def train_llama3():
|
||||
for g, new_g in zip(grads, loss.gradient(*optim.params)):
|
||||
apply_grad(g, new_g.uop)
|
||||
|
||||
loss_acc.assign(loss_acc + loss.flatten().float())
|
||||
return loss_acc.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
|
||||
loss_cpu = loss.flatten().float().to("CPU")
|
||||
return loss_cpu.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
|
||||
|
||||
@TinyJit
|
||||
def optim_step():
|
||||
@@ -1491,10 +1490,9 @@ def train_llama3():
|
||||
|
||||
lr_cpu = optim.lr.float().to("CPU")
|
||||
grad_norm_cpu = grad_norm.float().to("CPU")
|
||||
loss_cpu = loss_acc.to("CPU")
|
||||
Tensor.realize(lr_cpu, grad_norm_cpu, loss_cpu, loss_acc.assign(0), *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)
|
||||
Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)
|
||||
|
||||
return lr_cpu, grad_norm_cpu, loss_cpu
|
||||
return lr_cpu, grad_norm_cpu
|
||||
|
||||
@TinyJit
|
||||
@Context(TRAINING=0)
|
||||
@@ -1549,8 +1547,8 @@ def train_llama3():
|
||||
st = time.perf_counter()
|
||||
|
||||
stopped = False
|
||||
data_time, dev_time = 0, 0
|
||||
for _ in range(accum_steps:=grad_acc if i >= 2 else 1):
|
||||
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:
|
||||
@@ -1558,15 +1556,16 @@ def train_llama3():
|
||||
break
|
||||
mst = time.perf_counter()
|
||||
data_time += mst - ist
|
||||
minibatch(tokens)
|
||||
losses.append(minibatch(tokens).item())
|
||||
dev_time += time.perf_counter() - mst
|
||||
if stopped: break
|
||||
|
||||
gt = time.perf_counter()
|
||||
ret = optim_step()
|
||||
lr, grad_norm, loss = ret[0].item(), ret[1].item(), ret[2].item() / accum_steps
|
||||
lr, grad_norm = ret[0].item(), ret[1].item()
|
||||
et = time.perf_counter()
|
||||
|
||||
loss = sum(losses) / len(losses)
|
||||
optim_time = et - gt
|
||||
dev_time += optim_time
|
||||
step_time = et - st
|
||||
|
||||
@@ -114,11 +114,6 @@ def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
|
||||
amax_x2:Tensor|None, next_amax_x2:Tensor|None,
|
||||
grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None,
|
||||
grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None):
|
||||
if FUSED_SILU_W13 and MXFP4:
|
||||
from extra.llama_kernels.swiglu import swiglu
|
||||
out, *ret = matmul(swiglu(x_w13), w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
|
||||
next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
|
||||
return out, ret
|
||||
if FUSED_SILU_W13 and not MXFP4:
|
||||
from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
|
||||
x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
|
||||
|
||||
@@ -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, asm_gemm, can_use_asm_gemm
|
||||
from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
|
||||
from extra.gemm.moe_gemm import grouped_mx_gemm
|
||||
from extra.gemm.moe_routing import route, dispatch, combine
|
||||
|
||||
@@ -182,12 +182,12 @@ class GPTOSS:
|
||||
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 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, window=self.sliding_window if sliding else 0)
|
||||
attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
|
||||
elif sliding:
|
||||
if sliding:
|
||||
attn = self._sliding_attention(xq, xk, xv, sinks)
|
||||
elif 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 = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
|
||||
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)
|
||||
@@ -263,11 +263,7 @@ 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)
|
||||
|
||||
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
|
||||
logits = self.norm(h) @ self.output.T
|
||||
return logits
|
||||
|
||||
def _get_pads(uop:UOp) -> list[UOp]:
|
||||
|
||||
+2
-2
@@ -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="float32"
|
||||
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
|
||||
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))
|
||||
|
||||
@@ -46,7 +46,7 @@ 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 FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
|
||||
export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
|
||||
if [ -z "$FULL_LAYERS" ]; then
|
||||
export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
|
||||
fi
|
||||
|
||||
+3
-3
@@ -1,8 +1,8 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
|
||||
export PATH="$ROCM_PATH/bin:$PATH"
|
||||
export PATH="/opt/rocm-7.1.1/bin:$PATH"
|
||||
export ROCM_PATH="/opt/rocm-7.1.1"
|
||||
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="float32"
|
||||
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
|
||||
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))
|
||||
|
||||
|
||||
@@ -122,8 +122,7 @@ 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"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)))
|
||||
arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
|
||||
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))))
|
||||
|
||||
|
||||
+7
-19
@@ -288,17 +288,15 @@ def amd_build_program(prg:UOp) -> UOp:
|
||||
|
||||
class AMDAllocator(HCQAllocator['AMDDevice']):
|
||||
def __init__(self, dev:AMDDevice):
|
||||
super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())
|
||||
super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_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_copy_queue)
|
||||
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 _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_unmap(self, buf:HCQ2Buffer): self.dev.iface.unmap(buf)
|
||||
|
||||
@dataclass
|
||||
class AMDQueueDesc:
|
||||
ring: Buffer; read_ptr: Buffer; write_ptr: Buffer; doorbell: Buffer; put_value: Buffer # noqa: E702
|
||||
@@ -390,24 +388,15 @@ 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():
|
||||
mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
|
||||
mapped._owns_kfd_handle = True
|
||||
return mapped
|
||||
if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
|
||||
|
||||
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)
|
||||
@@ -479,7 +468,6 @@ 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
|
||||
@@ -593,7 +581,7 @@ class AMDDevice(HCQ2Compiled):
|
||||
|
||||
self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
|
||||
self.sdma_queues:dict = {}
|
||||
self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
|
||||
self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this
|
||||
|
||||
super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
|
||||
|
||||
|
||||
@@ -1,49 +0,0 @@
|
||||
import functools, math
|
||||
from tinygrad import Tensor, dtypes
|
||||
from tinygrad.uop.ops import UOp, KernelInfo
|
||||
from tinygrad.renderer import Estimates
|
||||
from extra.llama_kernels import alloc_like
|
||||
|
||||
LOG2E = 1.4426950408889634
|
||||
|
||||
@functools.cache
|
||||
def _custom_swiglu(out:UOp, x_w13:UOp) -> UOp:
|
||||
rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2
|
||||
n_elems = rows * hidden
|
||||
out, x_w13 = out.reshape(n_elems), x_w13.reshape(rows, 2*hidden)
|
||||
i = UOp.range(n_elems, 0)
|
||||
row, col = i // hidden, i % hidden
|
||||
act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float)
|
||||
sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal()
|
||||
store = out[i].store((act * sigmoid * gate).cast(out.dtype))
|
||||
return store.end(i).sink(arg=KernelInfo(f"swiglu_fwd_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=6*n_elems)))
|
||||
|
||||
@functools.cache
|
||||
def _custom_swiglu_bwd(grad_out:UOp, x_w13:UOp, grad_act:UOp) -> UOp:
|
||||
rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2
|
||||
n_elems = rows * hidden
|
||||
grad_out, x_w13, grad_act = grad_out.reshape(rows, 2*hidden), x_w13.reshape(rows, 2*hidden), grad_act.reshape(n_elems)
|
||||
i = UOp.range(n_elems, 0)
|
||||
row, col = i // hidden, i % hidden
|
||||
act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float)
|
||||
grad = grad_act[i].cast(dtypes.float)
|
||||
sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal()
|
||||
silu = act * sigmoid
|
||||
dact = grad_out[row, col].store((grad * (sigmoid + silu * (1.0 - sigmoid)) * gate).cast(grad_out.dtype))
|
||||
dgate = grad_out.after(dact)[row, hidden+col].store((grad * silu).cast(grad_out.dtype))
|
||||
return dgate.end(i).sink(arg=KernelInfo(f"swiglu_bwd_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=10*n_elems)))
|
||||
|
||||
def _swiglu_bwd(gradient:UOp, kernel:UOp):
|
||||
_, x_w13 = kernel.src[1:]
|
||||
axis = x_w13.axis if isinstance(x_w13.device, tuple) else None
|
||||
grad_out = alloc_like(x_w13.shape, dtypes.bfloat16, x_w13.device, axis)
|
||||
grad_out, *_ = Tensor.custom_kernel(grad_out, Tensor(x_w13, device=x_w13.device), Tensor(gradient, device=x_w13.device),
|
||||
fxn=_custom_swiglu_bwd)
|
||||
return (None, grad_out.uop)
|
||||
|
||||
def swiglu(x_w13:Tensor) -> Tensor:
|
||||
assert x_w13.dtype == dtypes.bfloat16 and x_w13.ndim >= 2 and x_w13.shape[-1] % 32 == 0
|
||||
*prefix, two_k = x_w13.shape
|
||||
axis = x_w13.uop.axis if isinstance(x_w13.device, tuple) else None
|
||||
out = alloc_like((*prefix, two_k//2), dtypes.bfloat16, x_w13.device, axis)
|
||||
return Tensor.custom_kernel(out, x_w13, fxn=_custom_swiglu, grad_fxn=_swiglu_bwd)[0]
|
||||
@@ -1,221 +0,0 @@
|
||||
# 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 numpy tqdm wandb tiktoken sentencepiece
|
||||
```
|
||||
|
||||
### 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 Reload amdgpu driver
|
||||
tinygrad's HCQ backend needs `/dev/kfd` which is created by the amdgpu kernel driver.
|
||||
If the driver was unloaded, reload it:
|
||||
```bash
|
||||
modprobe amdgpu
|
||||
ls /dev/kfd # should exist
|
||||
```
|
||||
|
||||
## 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
|
||||
```
|
||||
|
||||
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`
|
||||
|
||||
## Phase 4: wandb Login
|
||||
```bash
|
||||
wandb login
|
||||
```
|
||||
Enter API key from https://wandb.ai/authorize
|
||||
|
||||
## Phase 5: Run Training
|
||||
|
||||
### 5.1 Smoke test (beam search, 2 layers, real data)
|
||||
Always run beam first to validate the pipeline:
|
||||
```bash
|
||||
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=AMD:HIP \
|
||||
ROCM_PATH=/opt/rocm \
|
||||
bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
|
||||
```
|
||||
|
||||
### 5.2 Full training run
|
||||
```bash
|
||||
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=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
|
||||
```
|
||||
|
||||
## 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` | `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, direct KFD ioctl |
|
||||
| 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`.
|
||||
|
||||
### PCI device ID
|
||||
`lspci -v` shows device ID `0x75b0` and subsystem ID `0x75a0`:
|
||||
```
|
||||
83:00.0 Processing accelerators: ... Device 75b0
|
||||
Subsystem: ... Device 75a0
|
||||
```
|
||||
tinygrad's `PCIIface` in `ops_amd.py` and `hive_reset.py` did not list `0x75b0`, so the GPU was not found. Adding `0x75b0` to the device ID list in both files fixes the detection.
|
||||
|
||||
### amdgpu driver behavior
|
||||
On first boot, amdgpu loaded and bound to all 8 GPUs. On one boot it failed to initialize:
|
||||
```
|
||||
[ 799.780369] amdgpu 0000:83:00.0: Failed to alloc msi vectors
|
||||
[ 799.781476] amdgpu 0000:83:00.0: sw_init of IP block <vega20_ih> failed -22
|
||||
[ 799.782724] amdgpu 0000:83:00.0: amdgpu_device_ip_init failed
|
||||
[ 799.793885] amdgpu 0000:83:00.0: Fatal error during GPU init
|
||||
```
|
||||
On a subsequent boot, amdgpu initialized successfully (SMU initialized, VRAM ready). After unbinding all 8 GPUs from amdgpu, `rmmod amdgpu` wedged the module (stuck in "Unloading" state in `/proc/modules`), requiring a full VM reboot.
|
||||
|
||||
### `/dev/kfd`
|
||||
`/dev/kfd` exists when amdgpu is loaded. Opening it returns `OSError: [Errno 22] Invalid argument`.
|
||||
|
||||
### VRAM BAR reads all 0xFF
|
||||
After amdgpu initializes the GPU and is then unbound, reading the VRAM BAR (via `/sys/bus/pci/devices/0000:83:00.0/resource0`) returns all `0xFF` at all offsets — including the discovery table at `vram_size - 64KB`. tinygrad's `AMDev._run_discovery()` fails with `AssertionError: discovery signatures mismatch`.
|
||||
|
||||
A PCI reset (`echo 1 > /sys/bus/pci/devices/0000:83:00.0/reset`) did not change the VRAM contents — still all `0xFF`.
|
||||
|
||||
VRAM was also all `0xFF` when read via `/dev/mem` at the BAR physical address (`0xa0000000000`).
|
||||
|
||||
### VFIO attempt
|
||||
Bound the GPU to `vfio-pci` with `enable_unsafe_noiommu_mode=1`. The GPU bound successfully and `/dev/vfio/noiommu-0` appeared. Running tinygrad with `VFIO=1` still failed with the same `discovery signatures mismatch` — VRAM BAR still reads all `0xFF`.
|
||||
|
||||
### No IOMMU in guest
|
||||
`dmesg` has no `AMD-Vi` entries. PCI devices have no `iommu_group` symlink.
|
||||
|
||||
### 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.
|
||||
|
||||
### Current status: NOT WORKING
|
||||
tinygrad's `PCIIface` finds the GPU (after adding `0x75b0`) but `AMDev._run_discovery()` fails because the VRAM discovery table reads all `0xFF`. This was observed with the GPU unbound from any driver, after PCI reset, and with VFIO bound.
|
||||
+9
-55
@@ -110,49 +110,7 @@ def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
|
||||
return _sharded_empty(ref.shape, ref, axis)
|
||||
|
||||
@functools.cache
|
||||
def _windowed_lse(xq:Tensor, xk:Tensor, sinks, W:int) -> Tensor:
|
||||
B, N, H, hd = xq.shape
|
||||
H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
|
||||
q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
|
||||
k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
|
||||
k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
|
||||
sc_d = (q @ k.transpose(-1, -2)) * sm
|
||||
sc_p = (q @ k_prev.transpose(-1, -2)) * sm
|
||||
li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
|
||||
pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
|
||||
sc_d = (lj <= li).where(sc_d, -float("inf"))
|
||||
sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
|
||||
m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
|
||||
if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
|
||||
denom = (sc_d - m).exp().sum(-1, keepdim=True) + (sc_p - m).exp().sum(-1, keepdim=True)
|
||||
if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
|
||||
return (m + denom.log()).reshape(B, H, N).unsqueeze(2) # (B, H, 1, N), matches saved l_vec
|
||||
|
||||
def _windowed_delta(xq:Tensor, xk:Tensor, xv:Tensor, do:Tensor, sinks, W:int) -> Tensor:
|
||||
B, N, H, hd = xq.shape
|
||||
H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
|
||||
q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
|
||||
k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
|
||||
v = xv.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
|
||||
dob = do.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
|
||||
k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
|
||||
v_prev = v.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
|
||||
sc_d = (q @ k.transpose(-1, -2)) * sm
|
||||
sc_p = (q @ k_prev.transpose(-1, -2)) * sm
|
||||
li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
|
||||
pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
|
||||
sc_d = (lj <= li).where(sc_d, -float("inf"))
|
||||
sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
|
||||
m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
|
||||
if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
|
||||
e_d, e_p = (sc_d - m).exp(), (sc_p - m).exp()
|
||||
denom = e_d.sum(-1, keepdim=True) + e_p.sum(-1, keepdim=True)
|
||||
if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
|
||||
o = ((e_d / denom) @ v) + ((e_p / denom) @ v_prev)
|
||||
delta = (dob * o).sum(-1)
|
||||
return delta.reshape(B, H, N).unsqueeze(2)
|
||||
|
||||
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 _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 grad(dou:UOp, ker:UOp) -> tuple:
|
||||
do = Tensor(dou, device=dou.device)
|
||||
attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
|
||||
@@ -160,8 +118,6 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
|
||||
xq = Tensor(ker.src[3], device=ker.src[3].device)
|
||||
xk = Tensor(ker.src[4], device=ker.src[4].device)
|
||||
xv = Tensor(ker.src[5], device=ker.src[5].device)
|
||||
if window:
|
||||
l_vec = _windowed_lse(xq, xk, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
|
||||
|
||||
dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
|
||||
GROUP_SIZE = H_local // H_KV_local
|
||||
@@ -172,10 +128,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
|
||||
# delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
|
||||
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]
|
||||
if window:
|
||||
delta_vec = _windowed_delta(xq, xk, xv, do, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
|
||||
|
||||
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]
|
||||
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]
|
||||
|
||||
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)
|
||||
@@ -195,7 +149,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, window:int=0):
|
||||
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None):
|
||||
assert attn_mask is None, "attn_mask not supported"
|
||||
assert is_causal, "only causal attention supported"
|
||||
|
||||
@@ -222,18 +176,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, window=window)
|
||||
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)
|
||||
|
||||
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, window=window), 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), 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, window:int=0):
|
||||
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):
|
||||
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"-DWINDOW={window}"]
|
||||
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)}"]
|
||||
|
||||
Q_BLOCK_SIZE = 32
|
||||
NUM_WARPS = 8
|
||||
@@ -293,10 +247,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, window:int=0):
|
||||
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):
|
||||
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"-DWINDOW={window}"]
|
||||
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]
|
||||
|
||||
BLOCK_SIZE_KV = 256
|
||||
GROUP_SIZE = H // H_KV
|
||||
|
||||
@@ -422,8 +422,9 @@ class TestCustomKernel(unittest.TestCase):
|
||||
return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
|
||||
GlobalCounters.reset()
|
||||
y = run(x[0]).realize()
|
||||
# backends that support contiguous views don't launch extra kernels
|
||||
assert_kernel_count(2 if x[0].uop.contiguous_view() is None else 1)
|
||||
# it's copying the input and the output
|
||||
# TODO: subbuffer usage has runtime specific behavior, this will be fixed after the removal of SLICE.
|
||||
assert_kernel_count(2 if y.device in ("CL", "WEBGPU") else 1)
|
||||
self.assertEqual(y.tolist(), [1, 2, 3, 4])
|
||||
|
||||
@Context(DEV="CPU")
|
||||
|
||||
@@ -437,7 +437,7 @@ def reset_bufs(bufs:list[Buffer]):
|
||||
for buf in bufs: buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=memoryview(bytearray(buf.nbytes))))
|
||||
|
||||
def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
|
||||
apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[], check_default_opt=True):
|
||||
apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
|
||||
outbufs = real_bufs[:len(realized_ast.src)]
|
||||
wanna_output = [np.array(x).flatten() for x in wanna_output]
|
||||
buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs]
|
||||
@@ -459,7 +459,9 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
|
||||
for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
|
||||
|
||||
# Check correctness of handcoded optimiztions.
|
||||
if check_default_opt: check_opt(None)
|
||||
reset_bufs(outbufs)
|
||||
run_prg(opts=None)
|
||||
for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
|
||||
for x in opts: # Check custom transformations if any.
|
||||
check_opt(([Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1))] if apply_tc else [])+x)
|
||||
|
||||
|
||||
@@ -5,7 +5,6 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
|
||||
from extra.llama_kernels.fused_ce import fused_ce_loss
|
||||
from extra.llama_kernels import local_abs_max
|
||||
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
|
||||
from extra.llama_kernels.swiglu import swiglu
|
||||
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
|
||||
from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
|
||||
from test.helpers import needs_second_gpu, assert_kernel_count
|
||||
@@ -162,31 +161,5 @@ class TestFusedQKVRoPE(unittest.TestCase):
|
||||
ref = Tensor.cat(dq_ref, dk_ref, dv_ref, dim=3).reshape(*dx.shape).realize()
|
||||
with Context(DEBUG=0): self.assertTrue(dx.allclose(ref, atol=2e-2, rtol=2e-2).item(), "backward mismatch")
|
||||
|
||||
def run_swiglu(test:unittest.TestCase, shape:tuple[int, ...]) -> None:
|
||||
Tensor.manual_seed(0)
|
||||
x = (Tensor.randn(*shape) * 2).cast(dtypes.bfloat16).realize()
|
||||
hidden = x.shape[-1] // 2
|
||||
out, ref = swiglu(x), x[..., :hidden].silu() * x[..., hidden:]
|
||||
Tensor.realize(out, ref)
|
||||
with Context(DEBUG=0): test.assertTrue(out.allclose(ref, atol=2.5e-1, rtol=3e-2).item(), "SwiGLU forward mismatch")
|
||||
|
||||
grad = (Tensor.randn(*out.shape) * 2).cast(dtypes.bfloat16).realize()
|
||||
grad_x, grad_ref = out.gradient(x, gradient=grad)[0], ref.gradient(x, gradient=grad)[0]
|
||||
Tensor.realize(grad_x, grad_ref)
|
||||
test.assertEqual(grad_x.shape, shape)
|
||||
test.assertEqual(grad_x.dtype, dtypes.bfloat16)
|
||||
with Context(DEBUG=0): test.assertTrue(grad_x.allclose(grad_ref, atol=2.5e-1, rtol=3e-2).item(), "SwiGLU backward mismatch")
|
||||
|
||||
class TestSwiGLU(unittest.TestCase):
|
||||
def setUp(self):
|
||||
if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("need bfloat16")
|
||||
|
||||
def test_simple(self): run_swiglu(self, (2, 32, 64))
|
||||
|
||||
def test_llama_shape(self):
|
||||
if Device.DEFAULT != "AMD" or not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"):
|
||||
self.skipTest("only run on real machine for speed")
|
||||
run_swiglu(self, (2, 8192, 28672))
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
@@ -1,14 +0,0 @@
|
||||
import unittest, numpy as np
|
||||
from unittest.mock import patch
|
||||
from tinygrad import Device, Tensor
|
||||
from tinygrad.helpers import getenv
|
||||
from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in
|
||||
|
||||
@unittest.skipUnless(getenv("HCQ2") and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
|
||||
class TestHCQ2(unittest.TestCase):
|
||||
def test_copy_without_copy_queue(self):
|
||||
with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
|
||||
np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
-17
@@ -1,17 +0,0 @@
|
||||
from tinygrad import Device, Tensor, TinyJit, dtypes
|
||||
from tinygrad.helpers import Timing, Context
|
||||
|
||||
GPUS, DEPTH, SZ = 8, 4, 128 * 2**20
|
||||
WARMUP, ITERS = 3, 5
|
||||
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
|
||||
bufs = tuple(Tensor.empty(SZ, dtype=dtypes.uint8, device=dev).contiguous().realize() for _ in range(DEPTH) for dev in devs)
|
||||
|
||||
@TinyJit
|
||||
def all_to_all(*srcs:Tensor): return Tensor.realize(*(src.to(dst) for i,src in enumerate(srcs) for j,dst in enumerate(devs) if i % GPUS != j))
|
||||
|
||||
if __name__ == "__main__":
|
||||
with Context(ALL2ALL=1, JIT_BATCH_SIZE=0):
|
||||
for i in range(-WARMUP, ITERS):
|
||||
with Timing("ALL2ALL ", lambda ns: f" {SZ*GPUS*(GPUS-1)*DEPTH/ns:.2f} GB/s", enabled=i>=0):
|
||||
all_to_all(*bufs)
|
||||
for dev in devs: Device[dev].synchronize()
|
||||
+2
-2
@@ -48,8 +48,6 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
|
||||
else:
|
||||
assert isinstance(t, UOp), f"can't schedule {t}"
|
||||
linear, var_vals = Tensor(t).linear_with_vars()
|
||||
# test compiling the linear
|
||||
compile_linear(linear)
|
||||
kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1)
|
||||
for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink)
|
||||
if kernel_cnt != allowed:
|
||||
@@ -59,6 +57,8 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
|
||||
print("kernel", i+1)
|
||||
print(call.src[0])
|
||||
raise KernelCountException(allowed, kernel_cnt)
|
||||
# test compiling the linear
|
||||
compile_linear(linear)
|
||||
return linear, var_vals
|
||||
|
||||
def assert_kernel_count(expected:int):
|
||||
|
||||
@@ -74,7 +74,6 @@ class TestWhisper(unittest.TestCase):
|
||||
err
|
||||
)
|
||||
|
||||
@slow
|
||||
def test_transcribe_file1(self):
|
||||
self.assertEqual(transcribe_file(self.model, self.enc, TEST_FILE_1), TRANSCRIPTION_1)
|
||||
|
||||
@@ -90,7 +89,6 @@ class TestWhisper(unittest.TestCase):
|
||||
self.assertEqual(TRANSCRIPTION_1, transcriptions[0])
|
||||
self.assertEqual(TRANSCRIPTION_2, transcriptions[1])
|
||||
|
||||
@slow
|
||||
def test_transcribe_batch21(self):
|
||||
waveforms = [load_file_waveform(TEST_FILE_2), load_file_waveform(TEST_FILE_1)]
|
||||
transcriptions = transcribe_waveform(self.model, self.enc, waveforms)
|
||||
|
||||
@@ -27,15 +27,10 @@ def _make_linear(buffer_lists, copies=None):
|
||||
calls.append(UOp(Ops.CALL, src=(src0, *bufs)))
|
||||
return UOp(Ops.LINEAR, src=tuple(calls))
|
||||
|
||||
def _get_planned_view(buf:UOp) -> tuple[UOp, int, int]|None:
|
||||
view = buf.src[0] if buf.op is Ops.BITCAST else buf
|
||||
if view.op is not Ops.SHRINK or view.src[0].op is not Ops.BUFFER: return None
|
||||
return (arena:=view.src[0]), view.src[1].val * arena.dtype.itemsize, view.src[2].val * arena.dtype.itemsize
|
||||
|
||||
def _get_arena(buf, linear, result):
|
||||
for orig_si, new_si in zip(linear.src, result.src):
|
||||
for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
|
||||
if orig is buf and (planned:=_get_planned_view(new)) is not None: return planned[0]
|
||||
if orig is buf and new.op is Ops.SLICE: return new.src[0]
|
||||
return None
|
||||
|
||||
def check_assign(buffer_lists, copies=None):
|
||||
@@ -46,8 +41,8 @@ def check_assign(buffer_lists, copies=None):
|
||||
replace_map: dict[int, tuple[UOp, int, int]] = {}
|
||||
for orig_si, new_si in zip(linear.src, result.src):
|
||||
for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
|
||||
if (planned:=_get_planned_view(new)) is not None and id(orig) not in replace_map:
|
||||
replace_map[id(orig)] = planned
|
||||
if new.op is Ops.SLICE and id(orig) not in replace_map:
|
||||
replace_map[id(orig)] = (new.src[0], new.src[1].val * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
|
||||
|
||||
# verify pinned buffers are not planned
|
||||
for buf in held_bufs:
|
||||
|
||||
@@ -948,11 +948,6 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
|
||||
self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")
|
||||
|
||||
def test_equivalent_const_max(self):
|
||||
x = Variable("x", -10, 10)
|
||||
self.helper_test_variable((x < 0).where(0, x), 0, 10, "x.maximum(0)")
|
||||
self.helper_test_variable((0 < x).where(x, 0), 0, 10, "x.maximum(0)")
|
||||
|
||||
def test_where_combine(self):
|
||||
cond = Variable("x", 0, 3) < 2
|
||||
a = Variable("a", 0, 3)
|
||||
|
||||
@@ -239,15 +239,6 @@ class TestKernelOpts(unittest.TestCase):
|
||||
helper_linearizer_opt(a.sum().exp(), [[Opt(OptOps.PADTO, 0, 32)],])
|
||||
helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])
|
||||
|
||||
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
|
||||
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
|
||||
@unittest.expectedFailure
|
||||
def test_padto_group_full_unroll_sum(self):
|
||||
a = Tensor.ones(2, 28, 4096, dtype=dtypes.bfloat16).realize()
|
||||
out = ((a * 0.5).float().square()).sum(axis=(0, 2))
|
||||
opts_to_apply = [Opt(OptOps.GROUPTOP, 1, 256), Opt(OptOps.PADTO, 3, 32), Opt(OptOps.UNROLL, 2, 0), Opt(OptOps.UPCAST, 0, 7)]
|
||||
helper_linearizer_opt(out, [opts_to_apply], check_default_opt=False)
|
||||
|
||||
def test_padto_sum(self):
|
||||
N = 18
|
||||
# NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
|
||||
|
||||
@@ -79,8 +79,7 @@ class TestTensorCores(unittest.TestCase):
|
||||
@unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
|
||||
def test_tensor_cores(self):
|
||||
for tc in Device[Device.DEFAULT].renderer.tensor_cores:
|
||||
with self.subTest(tc=tc):
|
||||
helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)
|
||||
helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)
|
||||
|
||||
@unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
|
||||
def test_tensor_cores_nested_reduce(self):
|
||||
@@ -186,10 +185,10 @@ class TestTensorCores(unittest.TestCase):
|
||||
# skip fp8 tcs: the unoptimized ALU baseline quantizes products to fp8 (JAX promotion), which legitimately
|
||||
# differs from the MFMA path (f32 accumulation), so the baseline-vs-TC numerical gate can't hold for fp8.
|
||||
tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in not in dtypes.fp8s)
|
||||
x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
|
||||
x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
|
||||
r = x.matmul(y, dtype=tc.dtype_out)
|
||||
opts = [Opt(OptOps.UNROLL, 0, 2)]
|
||||
ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
|
||||
ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
|
||||
for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
|
||||
if u.op is Ops.WMMA:
|
||||
assert u.src[-1].src[0].op != Ops.STORE
|
||||
@@ -200,10 +199,10 @@ class TestTensorCores(unittest.TestCase):
|
||||
@unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation")
|
||||
def test_tensor_cores_unroll_casted_phi(self):
|
||||
tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
|
||||
x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
|
||||
x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
|
||||
r = x.matmul(y, dtype=tc.dtype_out)
|
||||
opts = [Opt(OptOps.UNROLL, 0, 2)]
|
||||
ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
|
||||
ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
|
||||
for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
|
||||
if u.op is Ops.WMMA:
|
||||
#assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
|
||||
@@ -216,10 +215,10 @@ class TestTensorCores(unittest.TestCase):
|
||||
def test_tensor_cores_unroll_casted_phi_with_children(self):
|
||||
# all STORE children are outside the loop
|
||||
tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
|
||||
x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
|
||||
x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
|
||||
r = x.matmul(y, dtype=tc.dtype_out).relu()
|
||||
opts = [Opt(OptOps.UNROLL, 0, 2)]
|
||||
ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
|
||||
ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
|
||||
for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
|
||||
if u.op is Ops.WMMA:
|
||||
#assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
import unittest
|
||||
|
||||
from tinygrad import Tensor, Device, Variable, dtypes
|
||||
from tinygrad.helpers import DEV, fetch, round_up
|
||||
from tinygrad.engine.realize import compile_linear
|
||||
from tinygrad.uop.ops import Ops
|
||||
from tinygrad import Tensor, Device, dtypes
|
||||
from tinygrad.helpers import fetch, round_up
|
||||
from extra.hevc.hevc import parse_hevc_file_headers, nv_gpu
|
||||
from extra.hevc.decode import hevc_decode
|
||||
|
||||
@@ -65,7 +63,7 @@ class TestHevc(unittest.TestCase):
|
||||
self.assertEqual(list(frame3.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
|
||||
self.assertEqual(list(frame3.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT == "NV" and not DEV.interface.startswith("MOCK"), "real NV only")
|
||||
@unittest.skipUnless(Device.DEFAULT == "NV", "NV only")
|
||||
def test_hevc_decode(self):
|
||||
url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
|
||||
dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
|
||||
@@ -85,22 +83,5 @@ class TestHevc(unittest.TestCase):
|
||||
self.assertEqual(f.dtype, dtypes.uint8)
|
||||
self.assertEqual(f.device, "NV")
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT == "NV", "NV only")
|
||||
def test_hevc_decode_compile(self):
|
||||
url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
|
||||
dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
|
||||
|
||||
opaque, frame_info, _, _, luma_w, luma_h, _ = parse_hevc_file_headers(dat)
|
||||
offset, sz, frame_pos, max_hist, _ = frame_info[1]
|
||||
out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
|
||||
history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)]
|
||||
decoded = Tensor(dat, device="NV")[offset:offset+sz].decode_hevc_frame(
|
||||
Variable("pos", 0, max_hist + 1).bind(frame_pos), out_image_size, opaque[1], history)
|
||||
|
||||
compiled = compile_linear(decoded.linear_with_vars()[0])
|
||||
self.assertTrue(any(call.src[0].op is Ops.PROGRAM for call in compiled.src))
|
||||
encdec_calls = [call for call in compiled.src if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "encdec"]
|
||||
self.assertEqual(len(encdec_calls), 1)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -13,18 +13,12 @@ V_TOKS = UOp.variable("toks", 1, 32) # 32 is the default chunk_size in generate
|
||||
class TestTransformerGenerate(unittest.TestCase):
|
||||
def test_warmup(self):
|
||||
model, calls = Transformer(TEST_CONFIG), []
|
||||
def generate(tokens, **kwargs):
|
||||
def generate(tokens):
|
||||
calls.append(tokens)
|
||||
yield from (1, 2)
|
||||
with patch.object(model, "generate", generate): model.warmup()
|
||||
self.assertEqual(calls, [[0], [0]])
|
||||
|
||||
def test_warmup_then_generate_with_default_chunk(self):
|
||||
# warmup must not capture JIT graphs that generate()'s default chunk_size then rejects
|
||||
model = Transformer(TEST_CONFIG)
|
||||
model.warmup()
|
||||
self.assertIsInstance(next(model.generate([5, 6, 7, 8])), int)
|
||||
|
||||
def test_first_recurrent_generate_before_state_init(self):
|
||||
model = Transformer(TEST_CONFIG)
|
||||
model.has_recurrent_block = True
|
||||
@@ -44,15 +38,6 @@ class TestTransformerGenerate(unittest.TestCase):
|
||||
next(model.generate([1, 2, 3, 4, 5, 42, 10]))
|
||||
self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])
|
||||
|
||||
def test_recurrent_divergent_prompt_restarts(self):
|
||||
model, calls = Transformer(TEST_CONFIG), []
|
||||
model.has_recurrent_block, model._cached_tokens = True, [1, 2, 9]
|
||||
def mock_call(self, tokens, start_pos, temperature):
|
||||
calls.append(start_pos)
|
||||
return Tensor([[42]])
|
||||
with patch.object(Transformer, '__call__', mock_call): next(model.generate([1, 2, 10, 11]))
|
||||
self.assertEqual(calls[0], V_START_POS.bind(0))
|
||||
|
||||
def test_template_starts_reasoning(self):
|
||||
router = StreamRouter(reasoning=True)
|
||||
self.assertEqual(list(router.route("reasoning</think>answer")),
|
||||
@@ -193,12 +178,6 @@ class TestTransformerGenerate(unittest.TestCase):
|
||||
# with temperature=2.0, we should see at least 2 distinct outputs across 5 runs
|
||||
self.assertGreater(len(runs), 1, "high temperature should produce varied outputs")
|
||||
|
||||
def test_recurrent_temperature_high_produces_variety(self):
|
||||
model = Transformer(TEST_CONFIG)
|
||||
model.has_recurrent_block = True
|
||||
outputs = {model.forward(Tensor([[1]]), 0, Tensor([2.0])).item() for _ in range(5)}
|
||||
self.assertGreater(len(outputs), 1)
|
||||
|
||||
def test_temperature_passed_to_forward(self):
|
||||
"""Temperature from generate should be passed through to __call__."""
|
||||
model = Transformer(TEST_CONFIG)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
from tinygrad import Device
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from tinygrad.dtype import dtypes
|
||||
@@ -10,27 +11,36 @@ class TestMetalGraph(unittest.TestCase):
|
||||
self.MetalGraph = MetalGraph
|
||||
self.dev = Device[Device.DEFAULT]
|
||||
|
||||
def metal_buf(self, offset, bitcast=False):
|
||||
size = 4 if bitcast else 1
|
||||
buf = UOp.new_buffer(Device.DEFAULT, offset+size, dtypes.uint8)
|
||||
if offset: buf = buf[offset:offset+size]
|
||||
return buf.bitcast(dtypes.float32) if bitcast else buf
|
||||
def metal_buf(self, offset):
|
||||
buf = MagicMock()
|
||||
if offset > 0:
|
||||
buf.op = Ops.SLICE
|
||||
src = MagicMock()
|
||||
src.dtype = dtypes.uint8
|
||||
buf.src = (src, UOp.const(offset))
|
||||
buf.dtype = dtypes.uint8
|
||||
else:
|
||||
buf.op = Ops.BUFFER
|
||||
buf.device = Device.DEFAULT
|
||||
return buf
|
||||
|
||||
def supports_uop(self, *bufs):
|
||||
return self.MetalGraph.supports_uop([self.dev], UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(*bufs))
|
||||
def call(self, *bufs):
|
||||
c = MagicMock()
|
||||
c.src = (MagicMock(op=Ops.PROGRAM),) + tuple(bufs)
|
||||
return c
|
||||
|
||||
def test_supports_uop_normal_offset(self):
|
||||
assert self.supports_uop(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF)) is True
|
||||
assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True
|
||||
|
||||
def test_supports_uop_overflow_offset(self):
|
||||
assert self.supports_uop(self.metal_buf(0), self.metal_buf(0x100000000)) is False
|
||||
assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False
|
||||
|
||||
def test_supports_uop_non_view_buf(self):
|
||||
assert self.supports_uop(self.metal_buf(0)) is True
|
||||
|
||||
def test_supports_uop_bitcast(self):
|
||||
assert self.supports_uop(self.metal_buf(0xFFFFFFFF, bitcast=True)) is True
|
||||
assert self.supports_uop(self.metal_buf(0x100000000, bitcast=True)) is False
|
||||
def test_supports_uop_nonmetal_buf(self):
|
||||
# non-SLICE ops should not be checked for offset
|
||||
buf = MagicMock()
|
||||
buf.op = Ops.BUFFER
|
||||
buf.device = Device.DEFAULT
|
||||
self.MetalGraph.supports_uop([self.dev], self.call(buf))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
import math, functools
|
||||
from dataclasses import dataclass
|
||||
from tinygrad.dtype import DType, dtypes
|
||||
from tinygrad.uop.ops import PatternMatcher, UOp, UPat, Ops
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x N)
|
||||
@@ -136,42 +135,6 @@ amd_cdna4 = amd_cdna_1616128 + amd_cdna_161632 + amd_cdna_161616
|
||||
|
||||
def get_amd(arch): return {"gfx942": amd_cdna3, "gfx950": amd_cdna4, "gfx1200": amd_rdna4, "gfx1201": amd_rdna4}.get(arch, amd_rdna3)
|
||||
|
||||
pm_validate_wmma_rdna3 = PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
|
||||
src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
|
||||
if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
|
||||
src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
|
||||
if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
|
||||
for j in range(x.max_numel()*2)))),
|
||||
arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
|
||||
for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
|
||||
(UPat(Ops.WMMA, name="x"), lambda x: x.replace(
|
||||
src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
|
||||
if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
|
||||
])
|
||||
|
||||
pm_validate_wmma_rdna4 = PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
|
||||
dtype=dtypes.uint16,
|
||||
src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
|
||||
.bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
|
||||
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
|
||||
if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
|
||||
])
|
||||
|
||||
pm_validate_wmma_cdna = PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
|
||||
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
|
||||
if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
|
||||
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
|
||||
if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
|
||||
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
|
||||
if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
|
||||
])
|
||||
# ***** Apple Metal *****
|
||||
|
||||
metal = [TensorCore(dims=(8,8,8), threads=32, elements_per_thread=(2,2,2), dtype_in=di, dtype_out=do,
|
||||
|
||||
@@ -336,8 +336,6 @@ class Compiled:
|
||||
pm_lower:Any = None
|
||||
pm_bufferize:Any = None
|
||||
|
||||
has_copy_queue:bool = True
|
||||
|
||||
def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):
|
||||
from tinygrad.renderer import Renderer
|
||||
self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
|
||||
|
||||
@@ -44,6 +44,8 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
|
||||
current_batch, current_batch_devs = [], []
|
||||
|
||||
for si in linear.src:
|
||||
if si.src[0].op is Ops.SLICE: continue
|
||||
|
||||
devs = dedup([Device[x] for b in si.src[1:] if b.op is not Ops.BIND for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
|
||||
graph_t = graph_class(devs[0]) if devs[0].graph is not None else None
|
||||
|
||||
@@ -178,7 +180,7 @@ class CapturedJit(Generic[ReturnType]):
|
||||
if call.op is not Ops.CALL: continue
|
||||
arg_uops = get_call_arg_uops(call)
|
||||
outs, ins = get_call_outs_ins(call)
|
||||
out |= {b for k in set(outs) - set(ins) if (b:=u if (cv:=(u:=arg_uops[k]).contiguous_view()) is None else cv[0]).op is Ops.BUFFER}
|
||||
out |= {arg_uops[k] for k in set(outs) - set(ins) if arg_uops[k].op in (Ops.BUFFER, Ops.SLICE)}
|
||||
return out
|
||||
|
||||
def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType:
|
||||
|
||||
+29
-21
@@ -4,8 +4,8 @@ import time, random, itertools, math, contextlib, weakref, array
|
||||
from dataclasses import dataclass, replace, field
|
||||
from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
|
||||
from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
|
||||
from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
|
||||
from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
|
||||
from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite
|
||||
from tinygrad.device import Device, Buffer, MultiBuffer
|
||||
from tinygrad.renderer import Estimates
|
||||
from tinygrad.codegen import to_program
|
||||
from tinygrad.codegen.opt.postrange import args_from_ast
|
||||
@@ -17,7 +17,7 @@ def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call
|
||||
def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
|
||||
ast = call.src[0]
|
||||
if ast.op is Ops.PROGRAM: return tuple(ast.arg.outs), tuple(ast.arg.ins)
|
||||
if ast.op is Ops.COPY: return (0,), (1,)
|
||||
if ast.op in (Ops.COPY, Ops.SLICE): return (0,), (1,)
|
||||
if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
|
||||
return (), ()
|
||||
|
||||
@@ -27,6 +27,9 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
|
||||
|
||||
ast, arg_uops = call.src[0], get_call_arg_uops(call)
|
||||
if ast.op is Ops.PROGRAM: return ast.arg.name
|
||||
if ast.op is Ops.SLICE:
|
||||
offset = ast.src[1].val * arg_uops[1].dtype.itemsize
|
||||
return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {offset:<10d}", "yellow")
|
||||
if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):7s}", "yellow")
|
||||
if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow")
|
||||
if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return colored(f"batched {len(ast.src[0].src)}", "cyan")
|
||||
@@ -137,7 +140,7 @@ class ExecContext:
|
||||
cache: bool = True
|
||||
|
||||
def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
|
||||
if b.op in (Ops.MSELECT, Ops.SHRINK) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
|
||||
if b.op in (Ops.SLICE, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
|
||||
if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
|
||||
return inputs[b.arg.slot] if b.op is Ops.PARAM else b
|
||||
def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
|
||||
@@ -151,6 +154,13 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d
|
||||
for x in call.src[0].toposort())
|
||||
for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {"_device_num": j} if has_dnum else {}
|
||||
|
||||
def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
|
||||
resolved = resolve_params(call, ctx.input_uops)
|
||||
bufs = [cast(Buffer, b.buffer) for b in resolved]
|
||||
bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].val*bufs[1].dtype.itemsize)
|
||||
with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv
|
||||
return None
|
||||
|
||||
def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
|
||||
for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
|
||||
dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
|
||||
@@ -200,30 +210,27 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
|
||||
return t[0]
|
||||
|
||||
def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
|
||||
if (info:=call.arg.aux).inputs is not None:
|
||||
if (inputs:=call.arg.aux.inputs) is not None:
|
||||
bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
|
||||
table = call.src[1+info.inputs].buffer
|
||||
table = call.src[1+inputs].buffer
|
||||
for j,dev in enumerate(call.arg.aux.device):
|
||||
addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
|
||||
mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q')
|
||||
wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected")
|
||||
mv[:len(addrs)] = addrs
|
||||
|
||||
exec_kernel(replace(ctx, update_stats=DEBUG>=3), call, ast)
|
||||
exec_kernel(replace(ctx, update_stats=False), call, ast)
|
||||
|
||||
tms = []
|
||||
for devices,name,estimates,prof in info.kernels:
|
||||
for device in devices:
|
||||
tm = None
|
||||
if prof:
|
||||
(d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
|
||||
if ctx.wait:
|
||||
d.synchronize(timeout=ctx.timeout)
|
||||
st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
|
||||
tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
|
||||
stat_call = call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates, kernels=())))
|
||||
with track_stats(ctx, stat_call, device, [], ctx.var_vals) as et: et[0] = tm
|
||||
return max(tms) if tms else None
|
||||
tms:list[float|None] = []
|
||||
for e in (aux:=call.arg.aux).prof: cast(Any, Device[e.device]).prof_ents[e.st_id] = e
|
||||
for d in [cast(Any, Device[x]) for x in aux.device]:
|
||||
with track_stats(ctx, call, d.device, [], ctx.var_vals) as et:
|
||||
if ctx.wait:
|
||||
d.synchronize(timeout=ctx.timeout)
|
||||
ts = [d.signal(i)._buf.cpu_view().view(fmt='Q')[0] for e in aux.prof if e.device == d.device for i in (e.st_id, e.en_id)]
|
||||
if ts: et[0] = float(max(ts)-min(ts))/d.timestamp_divider/1e6
|
||||
tms += et
|
||||
return tms[0]
|
||||
|
||||
# flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
|
||||
pm_flatten_linear = PatternMatcher([
|
||||
@@ -254,6 +261,7 @@ pm_optimize_local_size = PatternMatcher([
|
||||
])
|
||||
|
||||
pm_exec = PatternMatcher([
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.SLICE, name="ast"),), name="call", allow_any_len=True), exec_view),
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), exec_copy),
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="ast"),), name="call", allow_any_len=True), exec_kernel),
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="ast"),), name="call", allow_any_len=True), exec_encdec),
|
||||
@@ -268,7 +276,7 @@ def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:li
|
||||
if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
|
||||
if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
|
||||
linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
|
||||
if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
|
||||
if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE) if profile is None else profile)
|
||||
return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
|
||||
|
||||
def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear
|
||||
|
||||
+8
-8
@@ -1,9 +1,9 @@
|
||||
from __future__ import annotations
|
||||
import time
|
||||
START_TIME = time.perf_counter()
|
||||
import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
|
||||
import os, functools, platform, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
|
||||
from collections import defaultdict
|
||||
import shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
|
||||
import subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload
|
||||
|
||||
@@ -13,7 +13,8 @@ U = TypeVar("U")
|
||||
def prod(x:Iterable[T]) -> T|int: return functools.reduce(operator.mul, x, 1)
|
||||
|
||||
# NOTE: helpers is not allowed to import from anything else in tinygrad
|
||||
OSX, WIN = sys.platform == "darwin", sys.platform == "win32"
|
||||
OSX, WIN = platform.system() == "Darwin", sys.platform == "win32"
|
||||
ARCH_X86 = any(x in platform.processor() for x in ("Intel", "i386", "x86_64"))
|
||||
BASEDIR = pathlib.Path(__file__).parent
|
||||
|
||||
# fix colors on Windows, https://stackoverflow.com/questions/12492810/python-how-can-i-make-the-ansi-escape-codes-to-work-also-in-windows
|
||||
@@ -230,7 +231,7 @@ class _DEV(ContextVar):
|
||||
|
||||
DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
|
||||
IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
|
||||
JIT, JIT_BATCH_SIZE = ContextVar("JIT", 1), ContextVar("JIT_BATCH_SIZE", 32)
|
||||
JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
|
||||
CHUNK_SIZE = 2**20 # TinyFS content-addressed store: blob chunk + hash-tree node granularity
|
||||
WINO, CAPTURING, TRACEMETA, NO_COLOR = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1), ContextVar("NO_COLOR", 0)
|
||||
TRAINING = ContextVar("TRAINING", 0)
|
||||
@@ -453,9 +454,9 @@ def _ensure_downloads_dir() -> pathlib.Path:
|
||||
if pathlib.Path("/etc/tinybox-release").is_file():
|
||||
# try creating dir with sudo
|
||||
if not (downloads_dir := pathlib.Path("/raid/downloads")).exists():
|
||||
system(f"sudo mkdir -p {downloads_dir}")
|
||||
system(f"sudo chown tiny:root {downloads_dir}")
|
||||
system(f"sudo chmod 775 {downloads_dir}")
|
||||
subprocess.run(["sudo", "mkdir", "-p", downloads_dir], check=True)
|
||||
subprocess.run(["sudo", "chown", "tiny:root", downloads_dir], check=True)
|
||||
subprocess.run(["sudo", "chmod", "775", downloads_dir], check=True)
|
||||
return downloads_dir
|
||||
return pathlib.Path(cache_dir) / "downloads"
|
||||
|
||||
@@ -496,7 +497,6 @@ def fetch_fw(path:str, name:str, sha256:str) -> bytes:
|
||||
# *** Exec helpers
|
||||
|
||||
def system(cmd:str, **kwargs) -> str:
|
||||
import subprocess
|
||||
st = time.perf_counter()
|
||||
try: ret = subprocess.check_output(cmd.split(), stderr=subprocess.STDOUT, **kwargs).decode().strip()
|
||||
except subprocess.CalledProcessError as e:
|
||||
|
||||
+8
-13
@@ -1,7 +1,7 @@
|
||||
from __future__ import annotations
|
||||
import functools, itertools, pathlib
|
||||
from dataclasses import dataclass, replace
|
||||
from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
|
||||
from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
|
||||
from tinygrad.nn import Linear
|
||||
from tinygrad.llm.gguf import gguf_load
|
||||
from tinygrad.uop.ops import resolve
|
||||
@@ -187,8 +187,8 @@ class TransformerBlock(FFNBlock):
|
||||
|
||||
def _init_state(self, x:Tensor):
|
||||
if not hasattr(self, "cache_kv"):
|
||||
self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
|
||||
dtype=dtypes.default_float, device=x.device)
|
||||
# TODO: how is the dtype of this determined?
|
||||
self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim, device=x.device)
|
||||
self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
|
||||
|
||||
class MLATransformerBlock(FFNBlock):
|
||||
@@ -261,14 +261,13 @@ class GatedDeltaNetBlock(FFNBlock):
|
||||
def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
|
||||
B, T, _ = x.shape
|
||||
assert T == 1, "GatedDeltaNetBlock currently only supports T=1"
|
||||
is_kda = hasattr(self, "ssm_g_a")
|
||||
|
||||
# input processing
|
||||
x = x.half()
|
||||
out_gate = self.ssm_g_b(self.ssm_g_a(x)) if is_kda else self.attn_gate(x)
|
||||
out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
|
||||
out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
|
||||
beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
|
||||
alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x)
|
||||
alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
|
||||
alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
|
||||
self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
|
||||
|
||||
@@ -292,13 +291,14 @@ class GatedDeltaNetBlock(FFNBlock):
|
||||
|
||||
# output
|
||||
core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
|
||||
out_gate = out_gate.sigmoid() if is_kda else out_gate.silu()
|
||||
out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
|
||||
return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
|
||||
|
||||
# recurrent state can't be partially reused after divergence, force a full rebuild
|
||||
def _state_reset_ops(self):
|
||||
return [self.conv_state.assign(self.conv_state.const_like(0)),
|
||||
self.recurrent_state.assign(self.recurrent_state.const_like(0))] if hasattr(self, "conv_state") else []
|
||||
def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return 0 if prefix_len != cached_len else prefix_len
|
||||
|
||||
def _init_state(self, x):
|
||||
if not hasattr(self, "conv_state"):
|
||||
@@ -326,8 +326,7 @@ class Transformer:
|
||||
def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
|
||||
x = self.token_embd(tokens).float() # (B, T, D)
|
||||
for block in self.blk: x = block(x, start_pos)
|
||||
# only run the output projection on the last token
|
||||
logits = self.output(self.output_norm(x[:, -1:]))[:, -1, :]
|
||||
logits = self.output(self.output_norm(x))[:, -1, :]
|
||||
# Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp)
|
||||
return (logits / temperature.maximum(1e-12) - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True)
|
||||
|
||||
@@ -421,10 +420,6 @@ class Transformer:
|
||||
for _ in range(2): list(zip(range(2), self.generate([0])))
|
||||
|
||||
def get_start_pos(self, tokens:list[int]) -> int:
|
||||
# recurrent state can't be partially reused after divergence: reuse it only when tokens extend the cached prefix
|
||||
if self.has_recurrent_block:
|
||||
return len(self._cached_tokens) if self._cached_tokens and len(self._cached_tokens) < len(tokens) \
|
||||
and tokens[:len(self._cached_tokens)] == self._cached_tokens else 0
|
||||
prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
|
||||
return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import json, math, pathlib, struct, functools, io, zlib
|
||||
import json, math, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Callable, BinaryIO, Iterable, cast
|
||||
from tinygrad.tensor import Tensor
|
||||
@@ -219,7 +219,6 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr
|
||||
|
||||
@accept_filename
|
||||
def zip_extract(t: Tensor) -> dict[str, Tensor]:
|
||||
import zipfile
|
||||
files: dict[str, Tensor] = {}
|
||||
with zipfile.ZipFile(TensorIO(t), "r") as myzip:
|
||||
# sadly, the extra length needs to be read from the local header of each file.
|
||||
@@ -250,7 +249,6 @@ def tar_extract(t: Tensor) -> dict[str, Tensor]:
|
||||
tensors = nn.state.tar_extract(Tensor(pathlib.Path("archive.tar")))
|
||||
```
|
||||
"""
|
||||
import tarfile
|
||||
with tarfile.open(fileobj=TensorIO(t), mode="r") as tar:
|
||||
return {member.name:t[member.offset_data:member.offset_data+member.size] for member in tar if member.type == tarfile.REGTYPE}
|
||||
|
||||
@@ -305,7 +303,6 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:
|
||||
"FloatTensor": None, "Parameter": Parameter}
|
||||
whitelist = {"torch", "collections", "numpy", "_codecs"} # NOTE: this is not for security, only speed
|
||||
class Dummy: pass
|
||||
import pickle, zipfile, tarfile
|
||||
class TorchPickle(pickle.Unpickler):
|
||||
def find_class(self, module, name):
|
||||
module_root = module.split(".")[0]
|
||||
|
||||
@@ -279,9 +279,43 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n %trunc
|
||||
(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, rdna4=AMDLLVMRenderer.is_rdna4(target.arch), cdna=self.is_cdna:
|
||||
render_wmma_amd(ctx, wmma, cdna, rdna4))
|
||||
])
|
||||
if self.is_cdna: self.extra_matcher += tc.pm_validate_wmma_cdna
|
||||
if target.arch in {"gfx1100", "gfx1151"}: self.extra_matcher += tc.pm_validate_wmma_rdna3
|
||||
if target.arch in {"gfx1200", "gfx1201"}: self.extra_matcher += tc.pm_validate_wmma_rdna4
|
||||
if self.is_cdna:
|
||||
self.extra_matcher += PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
|
||||
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
|
||||
if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
|
||||
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
|
||||
if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
|
||||
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
|
||||
if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
|
||||
])
|
||||
if target.arch in {"gfx1100", "gfx1151"}:
|
||||
self.extra_matcher += PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
|
||||
src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
|
||||
if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
|
||||
src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
|
||||
if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
|
||||
for j in range(x.max_numel()*2)))),
|
||||
arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
|
||||
for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
|
||||
(UPat(Ops.WMMA, name="x"), lambda x: x.replace(
|
||||
src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
|
||||
if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
|
||||
])
|
||||
if target.arch in {"gfx1200", "gfx1201"}:
|
||||
self.extra_matcher += PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
|
||||
dtype=dtypes.uint16,
|
||||
src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
|
||||
.bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
|
||||
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
|
||||
if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
|
||||
])
|
||||
|
||||
def supported_dtypes(self): return {d for d in super().supported_dtypes()
|
||||
if (d not in dtypes.fp8_ocp or self.target.arch == "gfx950") and d not in dtypes.fp8_fnuz}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import glob, importlib, os, pathlib, subprocess
|
||||
import glob, importlib, os, pathlib, shutil, subprocess, tarfile, tempfile
|
||||
from tinygrad.helpers import fetch, flatten, system, getenv
|
||||
|
||||
root = (here:=pathlib.Path(__file__).parent).parents[2]
|
||||
@@ -31,7 +31,6 @@ def load(name, files, **kwargs):
|
||||
if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists() or getenv('REGEN'):
|
||||
files, kwargs['args'] = files() if callable(files) else files, args() if callable(args:=kwargs.get('args', [])) else args
|
||||
if (srcs:=kwargs.pop('srcs', None)):
|
||||
import tempfile, tarfile
|
||||
srcpath = (td:=tempfile.TemporaryDirectory(f"autogen-src-{name.replace('/','-')}")).name + "/"
|
||||
for src in (srcs if isinstance(srcs, list) else [srcs]):
|
||||
if 'tar' in src:
|
||||
@@ -158,7 +157,7 @@ def __getattr__(nm):
|
||||
*[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
|
||||
srcs="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
|
||||
dll=f"'tinymesa_cpu' if DEV.renderer == 'LVP' else 'tinymesa', {tinymesa_path}, emsg='pip install tinymesa==25.2.7.2'",
|
||||
prolog=["from tinygrad.helpers import DEV", "import gzip, base64, sysconfig, os"],
|
||||
prolog=["from tinygrad.helpers import DEV", "import gzip, base64, platform, sysconfig, os"],
|
||||
epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
|
||||
case "libclang":
|
||||
return load("libclang",
|
||||
|
||||
@@ -5,7 +5,7 @@ from typing import Literal, TypeAlias
|
||||
from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
|
||||
from tinygrad.runtime.support import c
|
||||
from tinygrad.helpers import DEV
|
||||
import gzip, base64, sysconfig, os
|
||||
import gzip, base64, platform, sysconfig, os
|
||||
dll = c.DLL('mesa', 'tinymesa_cpu' if DEV.renderer == 'LVP' else 'tinymesa', os.path.join(sysconfig.get_paths()['platlib'], 'tinymesa'), emsg='pip install tinymesa==25.2.7.2')
|
||||
class struct_u_printf_info(c.Struct): pass
|
||||
u_printf_info: TypeAlias = struct_u_printf_info
|
||||
|
||||
@@ -6,6 +6,7 @@ from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer, P
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import UOp, Ops, Variable
|
||||
from tinygrad.engine.jit import GraphRunner, MultiGraphRunner
|
||||
from tinygrad.runtime.ops_rdma import RDMACopyQueue
|
||||
|
||||
class HCQGraph(MultiGraphRunner):
|
||||
def __init__(self, *args, **kwargs):
|
||||
@@ -49,7 +50,7 @@ class HCQGraph(MultiGraphRunner):
|
||||
|
||||
self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: unwrap(dev.hw_compute_queue_t)() for dev in self.devices}
|
||||
self.copy_queues: dict[tuple[HCQCompiled, int], HWQueue] = {} # lazy allocation, keyed by (device, queue_idx)
|
||||
self.rdma_queues: dict[tuple[HCQCompiled, HCQCompiled], "RDMACopyQueue"] = {} # lazy allocation, keyed by device pair
|
||||
self.rdma_queues: dict[tuple[HCQCompiled, HCQCompiled], RDMACopyQueue] = {} # lazy allocation, keyed by device pair
|
||||
self.num_copy_queues: int = getenv("HCQ_NUM_SDMA", min(len(self.devices), 8) if ALL2ALL >= 1 else 1)
|
||||
self.num_rdma_ops: dict[tuple[HCQCompiled, HCQCompiled], int] = collections.defaultdict(int)
|
||||
|
||||
@@ -103,7 +104,6 @@ class HCQGraph(MultiGraphRunner):
|
||||
elif is_rdma:
|
||||
enqueue_queue = self.comp_queues[enqueue_dev]
|
||||
rdma_key = (cast(HCQCompiled, Device[bufs[0].device]).rdma_dev(), enqueue_dev.rdma_dev())
|
||||
from tinygrad.runtime.ops_rdma import RDMACopyQueue
|
||||
self.rdma_queues.setdefault(rdma_key, RDMACopyQueue(enqueue_dev.rdma_dev()))
|
||||
else:
|
||||
assert (enqueue_dev.hw_copy_queue_t is not None), "device must implement a copy queue"
|
||||
|
||||
@@ -113,6 +113,5 @@ class MetalGraph(GraphRunner):
|
||||
@staticmethod
|
||||
def supports_uop(batch_devs, new_call:UOp) -> bool:
|
||||
# Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range.
|
||||
for shrink in [s for src in new_call.src[1:] if (s:=src.src[0] if src.op is Ops.BITCAST else src).op is Ops.SHRINK]:
|
||||
if shrink.src[1].val * shrink.src[0].dtype.itemsize > 0xFFFFFFFF: return False
|
||||
if any(b.op is Ops.SLICE and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
|
||||
return GraphRunner.supports_uop(batch_devs, new_call)
|
||||
|
||||
@@ -790,14 +790,10 @@ class KFDIface:
|
||||
|
||||
def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
|
||||
xcc_id=0, idx=0):
|
||||
if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA and idx and self.kfd_ver >= (1,17):
|
||||
queue_type = kfd.KFD_IOC_QUEUE_TYPE_SDMA_BY_ENG_ID
|
||||
sdma_engine_id = idx % (self.props['num_sdma_engines'] + self.props.get('num_sdma_xgmi_engines', 0))
|
||||
else: sdma_engine_id = 0
|
||||
queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring.va_addr, ring_size=ring.size, gpu_id=self.gpu_id,
|
||||
queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
|
||||
eop_buffer_address=eop_buffer.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer.size if eop_buffer else 0, ctl_stack_size=ctl_stack_size,
|
||||
ctx_save_restore_address=cwsr_buffer.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size, sdma_engine_id=sdma_engine_id,
|
||||
ctx_save_restore_address=cwsr_buffer.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size,
|
||||
write_pointer_address=gart.va_addr+wptr, read_pointer_address=gart.va_addr+rptr+8*xcc_id)
|
||||
|
||||
if not hasattr(self, 'doorbells'):
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from __future__ import annotations
|
||||
import platform, sys, os, ctypes, ctypes.util, functools, mmap, threading, array, itertools
|
||||
import platform, sys, os, ctypes, functools, mmap, threading, array, itertools
|
||||
from dataclasses import replace
|
||||
from typing import cast
|
||||
from tinygrad.helpers import to_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le, partition
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import ctypes, struct, platform, pathlib, shutil, subprocess, sys, tarfile, tempfile
|
||||
import ctypes, struct
|
||||
from tinygrad.device import Compiler
|
||||
from tinygrad.helpers import DEBUG, system, fetch, unwrap
|
||||
from tinygrad.helpers import DEBUG, system
|
||||
from tinygrad.runtime.support.compiler_mesa import disas_adreno
|
||||
# see https://github.com/sirhcm/tinydreno
|
||||
from tinygrad.runtime.autogen import llvm_qcom
|
||||
@@ -10,17 +10,10 @@ def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
|
||||
class QCOMCompiler(Compiler):
|
||||
def __init__(self, arch:str):
|
||||
assert arch.split(',')[0] == "a630", "only a630 supported"
|
||||
if platform.machine() == "aarch64": self.arch, self.chip_id, self.llvm_inst = arch, 0x6030001, llvm_qcom.cl_compiler_create_llvm_instance()
|
||||
else:
|
||||
self.arch, self.chip_id, self.fs = arch, 0x6030001, tempfile.TemporaryDirectory()
|
||||
with tarfile.open(fetch('https://git.tinygrad.win/sirhcm/images/releases/download/v2/qcomcl.tar.gz')) as t: t.extractall(fs:=self.fs.name)
|
||||
if (qemu:=shutil.which("qemu-aarch64-static")): argv = f"{qemu} -cpu max,pauth=off -L {fs} {fs}/usr/bin/python3 {__file__} {arch}"
|
||||
else: argv = (f"docker run --rm -i --platform linux/aarch64 -v {fs}/usr:/usr -v {pathlib.Path(__file__).parents[2]}:/tinygrad "
|
||||
f"-e PYTHONPATH=/ -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3 /tinygrad/runtime/support/compiler_qcom.py {arch}")
|
||||
self.compiler_process = subprocess.Popen(argv.split(), stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
|
||||
self.arch, self.chip_id, self.llvm_inst = arch, 0x6030001, llvm_qcom.cl_compiler_create_llvm_instance()
|
||||
super().__init__(f"compile_qcomcl_{arch}")
|
||||
|
||||
def __del__(self): llvm_qcom.cl_compiler_destroy_llvm_instance(self.llvm_inst) if platform.machine() == "aarch64" else self.compiler_process.kill()
|
||||
def __del__(self): llvm_qcom.cl_compiler_destroy_llvm_instance(self.llvm_inst)
|
||||
|
||||
def __reduce__(self): return QCOMCompiler, (self.arch,)
|
||||
|
||||
@@ -32,10 +25,6 @@ class QCOMCompiler(Compiler):
|
||||
return handle
|
||||
|
||||
def compile(self, src) -> bytes:
|
||||
if platform.machine() != "aarch64":
|
||||
unwrap(self.compiler_process.stdin).write(struct.pack("I", len(src.encode())) + src.encode())
|
||||
if (lib:=unwrap(self.compiler_process.stdout).read(struct.unpack("I", unwrap(self.compiler_process.stdout).read(4))[0])): return lib
|
||||
raise RuntimeError("QCOM Compilation Error")
|
||||
ch = self.checked(llvm_qcom.cl_compiler_compile_source(self.llvm_inst, self.chip_id, llvm_qcom.CL_MODE_64BIT, b"", 0, 0, 0, src.encode(), 0,
|
||||
llvm_qcom.CL_SRC_STR, None))
|
||||
if DEBUG >= 8: print(system("llvm-dis", input=ctypes.string_at((comp:=ch.contents.compiled.contents).llvm_bitcode, comp.llvm_bitcode_size)))
|
||||
@@ -47,13 +36,3 @@ class QCOMCompiler(Compiler):
|
||||
return ret
|
||||
|
||||
def disassemble(self, lib: bytes): disas_adreno(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)], self.chip_id)
|
||||
|
||||
if __name__ == "__main__":
|
||||
compiler = QCOMCompiler(sys.argv[1])
|
||||
while (amt:=sys.stdin.buffer.read(4)):
|
||||
try: lib = compiler.compile(sys.stdin.buffer.read(struct.unpack("I", amt)[0]).decode())
|
||||
except Exception as e:
|
||||
lib = b""
|
||||
print(e, file=sys.stderr, flush=True)
|
||||
sys.stdout.buffer.write(struct.pack("I", len(lib)) + lib)
|
||||
sys.stdout.buffer.flush()
|
||||
|
||||
@@ -33,13 +33,13 @@ class HCQInfo:
|
||||
|
||||
input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
|
||||
inputs:int|None = None
|
||||
kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...]], ...] = ()
|
||||
prof:tuple[ProfileGraphEntry, ...] = () # st_id/en_id are timestamp signal slots until collect
|
||||
|
||||
def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
|
||||
|
||||
def unwrap_mstack(u):
|
||||
if u.op is Ops.MSTACK: return tuple(x for s in u.src for x in unwrap_mstack(s))
|
||||
return unwrap_mstack(u.src[0]) if u.op is Ops.MSELECT else (u,)
|
||||
return unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,)
|
||||
|
||||
def is_value_known_at_link(val:UOp) -> bool:
|
||||
runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)]
|
||||
@@ -96,24 +96,12 @@ pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_
|
||||
|
||||
def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS)
|
||||
|
||||
def hcq_call_devs(call:UOp) -> Any|None: return next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)
|
||||
|
||||
def stage_copy(dst:UOp, src:UOp) -> UOp|None:
|
||||
if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
|
||||
|
||||
stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.itemsize, dtypes.uint8)
|
||||
return UOp(Ops.LINEAR, src=(src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
|
||||
|
||||
def kernel_copy(call:UOp, dst:UOp, src:UOp) -> UOp|None:
|
||||
if (devs:=hcq_call_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None
|
||||
d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2))
|
||||
ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1)
|
||||
return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src))
|
||||
|
||||
pm_insert_copy_staging = PatternMatcher([
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy),
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call"), kernel_copy)
|
||||
])
|
||||
pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
|
||||
|
||||
# *****************
|
||||
# 2. deps
|
||||
@@ -141,14 +129,14 @@ def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]
|
||||
|
||||
# opt2: keep latest dep per (dep device, queue, cur lane)
|
||||
latest = {((dep[0][dlane], dep[1]), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0][2])}
|
||||
deps:dict[tuple, dict[int, list[int]]] = collections.defaultdict(lambda: collections.defaultdict(list))
|
||||
for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane].append(dlane)
|
||||
deps:dict[tuple, list[int|None]] = collections.defaultdict(lambda: [None]*len(devices))
|
||||
for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane] = dlane
|
||||
|
||||
waits = []
|
||||
for (ddevs, dqueue, dtag), by_lane in deps.items():
|
||||
for ls in itertools.zip_longest(*(by_lane[lane] for lane in range(len(devices)))):
|
||||
s = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue]) for dl, d in zip(ls, devices)])
|
||||
waits.append(UOp(Ops.INS, arg="wait", src=(s, UOp.const(dtag + 1, dtypes.uint64))))
|
||||
for (ddevs, dqueue, dtag), lanes in deps.items():
|
||||
sig = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue])
|
||||
for dl, d in zip(lanes, devices)])
|
||||
waits.append(UOp(Ops.INS, arg="wait", src=(sig, UOp.const(dtag + 1, dtypes.uint64))))
|
||||
return waits, {dtag for _, _, dtag in deps}
|
||||
|
||||
def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
|
||||
@@ -199,10 +187,10 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
|
||||
signal_tags |= cur_signal_tags
|
||||
|
||||
# build fences and finalizers
|
||||
fences, fins, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
|
||||
fences, finalizers, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
|
||||
signal_tags |= finalizer_signal_tags
|
||||
|
||||
src, kerns = [], []
|
||||
src, prof = [], []
|
||||
for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
|
||||
# first queue use, sync prior device work with the device timeline
|
||||
if batch_info.index((devices, queue)) == tag:
|
||||
@@ -212,24 +200,24 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
|
||||
# and make hcq call
|
||||
name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call))
|
||||
ts_ids = [next(UOp.unique_num) for _ in range(2)] if profile else []
|
||||
kerns.append((devices, name, info.estimates, tuple(ts_ids)))
|
||||
prof += [ProfileGraphEntry(d, name, *ts_ids) for d in devices if ts_ids]
|
||||
|
||||
ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_signal(devices, s),)) for s in ts_ids]
|
||||
q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:]
|
||||
|
||||
# signal the queue if someone waits for us
|
||||
if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
|
||||
src.append(make_call(f"submit {name}", make_submit(*q, devs=devices, queue=queue).sink(), info))
|
||||
src.append(make_call(name, make_submit(*q, devs=devices, queue=queue).sink(), info))
|
||||
|
||||
# append batch timestamps to finalizers
|
||||
fins = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, kernels=tuple(x for x in kerns if set(x[0]) & set(a.device))))) for f in fins]
|
||||
return fences + src + fins
|
||||
finalizers = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, prof=tuple(e for e in prof if e.device in a.device)))) for f in finalizers]
|
||||
return fences + src + finalizers
|
||||
|
||||
def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
|
||||
srcs:list[UOp] = []
|
||||
batch:list[tuple[UOp, tuple[str, ...]]] = []
|
||||
for call in l.src:
|
||||
if (devs:=hcq_call_devs(call)) is not None: batch.append((call, to_tuple(devs)))
|
||||
if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs)))
|
||||
else: srcs, batch = srcs + _finalize_batch(batch, profile) + [call], []
|
||||
return l.replace(src=tuple(srcs + _finalize_batch(batch, profile)))
|
||||
|
||||
@@ -377,15 +365,14 @@ pm_replace_params = PatternMatcher([
|
||||
|
||||
# *****************
|
||||
|
||||
def resolve_getaddr_view(bv:UOp, g:UOp) -> UOp:
|
||||
def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
|
||||
base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ())
|
||||
if bv.op is Ops.BITCAST: return UOp(Ops.GETADDR, src=(base,), arg=g.arg)
|
||||
itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
|
||||
itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
|
||||
return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64)
|
||||
|
||||
pm_early_simplify = PatternMatcher([
|
||||
(UPat(Ops.GETADDR, src=(UPat((Ops.SHRINK, Ops.BITCAST), name="bv").or_after(),), name="g"), resolve_getaddr_view),
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.SHRINK, name="bv"),), allow_any_len=True, name="x"),
|
||||
(UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"),), allow_any_len=True, name="x"),
|
||||
lambda bv,x: x.replace(src=(bv.src[0], x.src[1] + bv.src[1].cast(x.src[1].dtype), *x.src[2:]))),
|
||||
])
|
||||
|
||||
@@ -403,7 +390,7 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
|
||||
sizes[b.tag] = offs[b] + b.max_numel()
|
||||
counts = collections.Counter(b.tag for b in bufs)
|
||||
bases = {b.tag:UOp.placeholder((sizes[b.tag],), b.dtype, next(UOp.unique_num), device=b.device).rtag(b.tag) for b in bufs if counts[b.tag] > 1}
|
||||
subs = {b:bases[b.tag][(off:=offs.get(b, 0)):off+b.max_numel()] for b in bufs if b.tag in bases}
|
||||
subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
|
||||
return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
|
||||
pm_pack_placeholders = PatternMatcher([
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
|
||||
@@ -494,7 +481,7 @@ pm_resolve_patches = PatternMatcher([
|
||||
(UPat(name="buf").index(UPat(Ops.RANGE), allow_any_len=True)
|
||||
.store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast()).index(UPat(Ops.RANGE), allow_any_len=True).load())
|
||||
.end(UPat(Ops.RANGE)), fold_binary),
|
||||
(UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
|
||||
(UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
|
||||
])
|
||||
|
||||
pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: panic(RuntimeError, f"AFTER left at hcq_link: {a.src[0].op}"))])
|
||||
@@ -624,8 +611,6 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
|
||||
if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
|
||||
return self._do_map(buf)
|
||||
|
||||
def _do_unmap(self, mb): self.dev.iface.free(mb)
|
||||
|
||||
@suppress_finalizing
|
||||
def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None):
|
||||
if options is not None and options.external_ptr is not None: return
|
||||
@@ -634,6 +619,6 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
|
||||
|
||||
def _unmap(self, mb):
|
||||
self.dev.synchronize()
|
||||
self._do_unmap(mb)
|
||||
self.dev.iface.free(mb)
|
||||
|
||||
def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size)
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
from __future__ import annotations
|
||||
import os, mmap, array, functools, ctypes, ctypes.util, select, contextlib, dataclasses, sys, itertools, struct, socket
|
||||
import subprocess, time, enum, atexit
|
||||
import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, itertools, struct, socket, subprocess, time, enum, atexit
|
||||
from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, unwrap, fetch, system, _ensure_downloads_dir, DEBUG, flatten, pluralize
|
||||
from tinygrad.runtime.autogen import libc, pci, vfio
|
||||
from tinygrad.runtime.autogen import libc, pci, vfio, iokit, corefoundation
|
||||
from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer, hcq_filter_visible_devices
|
||||
from tinygrad.runtime.support.memory import VirtMapping, AddrSpace, BumpAllocator
|
||||
from tinygrad.runtime.support.usb import USB3, CustomASM24Controller, USBMMIOInterface
|
||||
@@ -56,7 +55,6 @@ class _System:
|
||||
def pci_scan_bus(self, vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None=None) -> list[str]:
|
||||
all_devs = []
|
||||
if OSX:
|
||||
from tinygrad.runtime.autogen import iokit, corefoundation
|
||||
def read_prop(svc, key) -> int:
|
||||
cfkey = corefoundation.CFStringCreateWithCString(None, key.encode(), corefoundation.kCFStringEncodingUTF8)
|
||||
cfdata = ctypes.cast(iokit.IORegistryEntryCreateCFProperty(svc, ctypes.cast(cfkey, iokit.CFStringRef), None, 0), corefoundation.CFDataRef)
|
||||
|
||||
@@ -23,7 +23,7 @@ class IndexingContext:
|
||||
return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(0)
|
||||
|
||||
|
||||
ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER,
|
||||
ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER, Ops.SLICE,
|
||||
Ops.CONST, Ops.BIND, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
|
||||
Ops.LOAD, Ops.CALL, Ops.FUNCTION}
|
||||
|
||||
@@ -34,6 +34,10 @@ def realize_srcs(ctx:IndexingContext, rb:UOp) -> None:
|
||||
if s.base.op not in ALWAYS_CONTIGUOUS: ctx.realize_map[s] = None
|
||||
|
||||
def realize_store_after_src(ctx:IndexingContext, dest:UOp, src:UOp):
|
||||
# don't realize SLICE when it's the direct source of STORE+AFTER — the target buffer is the output
|
||||
if src.op is Ops.SLICE and src in ctx.realize_map \
|
||||
and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
|
||||
del ctx.realize_map[src]
|
||||
# you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
|
||||
if dest.base in src.backward_slice_with_self: ctx.realize_map[src] = None
|
||||
|
||||
@@ -70,7 +74,7 @@ def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
|
||||
# TODO: srcs contain (real data srcs, something else, ranges) and the boundary is confusing. see range_start
|
||||
def data_srcs(op:Ops, src:tuple[UOp, ...]) -> tuple[UOp, ...]:
|
||||
if op in {Ops.PARAM, Ops.BUFFER, Ops.RANGE, Ops.SPECIAL, Ops.BIND}: return ()
|
||||
if op in GroupOp.Movement|{Ops.INDEX, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
|
||||
if op in GroupOp.Movement|{Ops.INDEX, Ops.SLICE, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
|
||||
return src
|
||||
|
||||
def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
|
||||
@@ -80,7 +84,7 @@ def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
|
||||
for i, s in enumerate(x.src):
|
||||
new_src = s
|
||||
src_rngs = broadcast_rngs(x, s, ctx.range_map[x][0]) if x in ctx.range_map else ()
|
||||
if s.op in {Ops.PARAM, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
|
||||
if s.op in {Ops.PARAM, Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
|
||||
if x in ctx.range_map and i < data_src_count: new_src = new_src.index(*src_rngs)
|
||||
elif s in ctx.realize_map:
|
||||
realized_ranges = ctx.realize_map[s]
|
||||
|
||||
@@ -52,9 +52,11 @@ def memory_plan_rewrite(linear:UOp, held_bufs:set[UOp]|None=None) -> UOp:
|
||||
peaks[_key(buf)] = (max(peaks[_key(buf)][0], offsets[buf] + buf.max_numel() * buf.dtype.itemsize), peaks[_key(buf)][1])
|
||||
arena_sizes = {key: round_up(peak, block_size) for key, (peak, _) in peaks.items()}
|
||||
|
||||
# build replace_map: each buffer becomes a SHRINK/BITCAST into a shared per-device-lane arena
|
||||
# build replace_map: each buffer becomes a SLICE into a shared per-device-lane arena
|
||||
arenas = {key: UOp.new_buffer(key[0], sz, dtypes.int8) for key, sz in arena_sizes.items()}
|
||||
replace_map = {buf_uop:arenas[_key(buf_uop)][offset:offset+buf_uop.nbytes()].bitcast(buf_uop.dtype) for buf_uop, offset in offsets.items()}
|
||||
replace_map:dict[UOp, UOp] = {}
|
||||
for buf_uop, offset in offsets.items():
|
||||
replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(offset)), buf_uop.max_numel())
|
||||
|
||||
if DEBUG >= 1 and (omem:=sum(nbytes.values()) / 1e6) != (nmem:=sum(arena_sizes.values()) / 1e6):
|
||||
print(f"memory reduced from {omem:.2f} MB -> {nmem:.2f} MB, {len(first_appearance)} -> {len(arenas)} bufs")
|
||||
|
||||
+25
-27
@@ -23,7 +23,6 @@ class AllocCtx:
|
||||
bases: set[UOp] = field(default_factory=set)
|
||||
assigns: list[UOp] = field(default_factory=list)
|
||||
replacements: list[UOp] = field(default_factory=list)
|
||||
views: set[UOp] = field(default_factory=set)
|
||||
|
||||
def tag_uop(ctx:AllocCtx, x:UOp):
|
||||
if x.tag is not None: return None
|
||||
@@ -64,37 +63,40 @@ def replace_contig_with_store_after(u:UOp):
|
||||
def replace_store_after_with_contig(u:UOp, src:UOp):
|
||||
assigned_to = u
|
||||
while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.UNSHARD}: assigned_to = assigned_to.src[0].base
|
||||
if assigned_to.op is not Ops.BUFFER: return src.contiguous(tag=u.tag)
|
||||
if assigned_to.op not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag)
|
||||
|
||||
def _make_buffer_view(src:UOp) -> UOp|None:
|
||||
if (cv := src.contiguous_view()) is None: return None
|
||||
(buf, offset), size = cv, src.max_numel() * src.element_size() // cv[0].element_size()
|
||||
if buf.op is not Ops.BUFFER: return None
|
||||
# NB: make offset a UOp.variable here to do the offset computation in the kernels
|
||||
return buf[offset:offset+size].bitcast(src.dtype)
|
||||
|
||||
def contiguous_mops_to_view(ctx:AllocCtx, c:UOp, src:UOp):
|
||||
"""MOPS(BUFFER) → SHRINK when movement ops collapse to a contiguous range."""
|
||||
"""If movement ops on src collapse to a contiguous range, return SLICE. Otherwise None."""
|
||||
if (offset := src.contiguous_view_offset()) is None: return None
|
||||
buf = src.base
|
||||
while buf.op is Ops.BITCAST: buf = buf.src[0].base
|
||||
if buf.op not in {Ops.BUFFER, Ops.UNSHARD}: return None
|
||||
if buf.op is Ops.SLICE:
|
||||
byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize
|
||||
buf = buf.src[0]
|
||||
if byte_offset % buf.dtype.itemsize != 0: return None
|
||||
offset = byte_offset // buf.dtype.itemsize
|
||||
return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(offset)), src.numel())
|
||||
|
||||
def contiguous_mops_to_view(c:UOp, src:UOp):
|
||||
"""MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range."""
|
||||
buf = src.base
|
||||
if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.UNSHARD}: return None
|
||||
if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE} and c.op is not Ops.BITCAST: return None
|
||||
if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None
|
||||
|
||||
# no symbolic shape
|
||||
if not all_int(c.shape): return None
|
||||
|
||||
if buf.op is not Ops.UNSHARD and (view := _make_buffer_view(src)) is not None:
|
||||
ctx.views.add(view)
|
||||
view = view.reshape(c.shape)
|
||||
return c.replace(src=(view,)+c.src[1:]) if c.op in {Ops.COPY, Ops.STORE} else view
|
||||
view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape)
|
||||
return c.replace(src=(view,)) if c.op is Ops.COPY else view
|
||||
|
||||
# for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SHRINK on the resolved result
|
||||
# for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result
|
||||
if not isinstance(c.device, str):
|
||||
from tinygrad.schedule.multi import multi_pm
|
||||
resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view")
|
||||
if resolved.op is not Ops.UNSHARD: return None
|
||||
if (view := _make_buffer_view(resolved.src[0])) is None: return None
|
||||
ctx.views.add(view)
|
||||
return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:])
|
||||
return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag)
|
||||
|
||||
return None
|
||||
|
||||
@@ -149,9 +151,8 @@ pm_early_transform_tensor_graph = PatternMatcher([
|
||||
# resolve TUPLE+GETTUPLE (for precompiled calls)
|
||||
(UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
|
||||
|
||||
# fold MOPS+BITCAST over BUFFER into SHRINK when movement ops collapse to contiguous range
|
||||
(UPat((Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BITCAST}, name="src"),), name="c"), contiguous_mops_to_view),
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.BITCAST, name="src"), UPat()), name="c", allow_any_len=True), contiguous_mops_to_view),
|
||||
# fold MOPS+BITCAST over BUFFER/SLICE into SLICE when movement ops collapse to contiguous range
|
||||
(UPat((Ops.BITCAST, Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BUFFER}, name="src"),), name="c"), contiguous_mops_to_view),
|
||||
|
||||
# remove contiguous on movement ops before a copy on disk
|
||||
(UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, name="copy"), lambda x,copy:
|
||||
@@ -200,8 +201,6 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
|
||||
return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
|
||||
addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
|
||||
|
||||
def replace_input_view(ctx:AllocCtx, b:UOp): return replace_input_buffer(ctx, b) if b in ctx.views else None
|
||||
|
||||
pm_finalize_call = PatternMatcher([
|
||||
(UPat(Ops.AFTER, name="x"), finalize_after),
|
||||
(UPat(Ops.COPY, name="x"), lambda ctx,x: ctx.assigns.append(x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
|
||||
@@ -211,9 +210,8 @@ pm_replace_buf = PatternMatcher([
|
||||
# replace BUFFER with PARAM for cache key normalization
|
||||
(UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
|
||||
replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
|
||||
# replace SHRINK with PARAM
|
||||
(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), name="b", allow_any_len=True), replace_input_view),
|
||||
(UPat(Ops.BITCAST, src=(UPat.any(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), allow_any_len=True), UPat(Ops.BUFFER)),), name="b"), replace_input_view),
|
||||
# replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
|
||||
(UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer),
|
||||
# strip value from BIND for cache key normalization, so different values hit same cache
|
||||
(UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer),
|
||||
])
|
||||
@@ -231,7 +229,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
|
||||
big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="number the uops")
|
||||
|
||||
# here we can break the tensor graph. this is the only place you need to maintain numbered tags
|
||||
big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, ctx=ctx, name="early transform tensor graph")
|
||||
big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph")
|
||||
|
||||
# here we construct the final buffer_map: as-built nodes -> their final storage. values are never keys
|
||||
graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call")
|
||||
|
||||
@@ -93,7 +93,7 @@ class Ops(FastEnum):
|
||||
CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()
|
||||
|
||||
# buffer ops
|
||||
STAGE = auto(); COPY = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()
|
||||
STAGE = auto(); COPY = auto(); SLICE = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()
|
||||
|
||||
# the core 6 movement ops! these only exist in the tensor graph
|
||||
RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); FLIP = auto()
|
||||
|
||||
+40
-27
@@ -45,7 +45,8 @@ axis_colors = {AxisType.DEVICE: "green", AxisType.GLOBAL: "blue", AxisType.THREA
|
||||
axis_to_pos = {AxisType.DEVICE: -2, AxisType.WEAK: -1, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1,
|
||||
AxisType.LOCAL: 2, AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
|
||||
|
||||
range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1, Ops.LINEAR: 0}
|
||||
range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
|
||||
Ops.SLICE: 2, Ops.LINEAR: 0}
|
||||
|
||||
# https://en.wikipedia.org/wiki/Identity_element
|
||||
def identity_element(op:Ops, dt:DType) -> PyConst: return dt.const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dt.min}[op])
|
||||
@@ -170,6 +171,9 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
|
||||
return arg.dtype
|
||||
case Ops.BINARY:
|
||||
return dtypes.uint8
|
||||
case Ops.SLICE:
|
||||
# TODO: slice just shouldn't exist
|
||||
return None
|
||||
case Ops.CAST | Ops.BITCAST:
|
||||
assert isinstance(arg, DType), f"CAST/BITCAST arg must be DType, got {arg}"
|
||||
return arg
|
||||
@@ -217,7 +221,7 @@ class UOpMetaClass(type):
|
||||
return created
|
||||
|
||||
# some uops map to other stuff
|
||||
buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/view uops to their device Buffers
|
||||
buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/SLICE uops to their device Buffers
|
||||
all_metadata:weakref.WeakKeyDictionary[UOp, tuple[Metadata, ...]] = weakref.WeakKeyDictionary() # TODO: should this be here?
|
||||
|
||||
# recursive_property replaces functools.cached_property in recursive UOp functions to prevent RecursionError
|
||||
@@ -382,6 +386,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
case Ops.BUFFER:
|
||||
if len(self.src): return self.src[0].as_shape
|
||||
return ()
|
||||
case Ops.SLICE:
|
||||
# HACK: SLICE is used inside kernels, so we set the shape to () if it's on an INDEX
|
||||
if self.src[0].op is Ops.INDEX: return ()
|
||||
return (self.arg,)
|
||||
case Ops.CUSTOM | Ops.CUSTOMI:
|
||||
if self.dtype is dtypes.void: return None
|
||||
input_shapes = [x._shape for x in self.src if x._shape is not None]
|
||||
@@ -560,9 +568,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
in_tuple = self.src[0] if self.op is Ops.FUNCTION else self
|
||||
assert in_tuple.op is Ops.TUPLE, f"gettuple requires FUNCTION or TUPLE source, got {self.op}"
|
||||
return UOp(Ops.GETTUPLE, src=(self,), arg=idx)
|
||||
def group(*srcs:UOp|None, **kwargs): # pylint: disable=no-self-argument
|
||||
def group(*srcs:UOp|None): # pylint: disable=no-self-argument
|
||||
if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
|
||||
return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]), **kwargs)
|
||||
return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]))
|
||||
def index(self, *srcs:UOp|int|None, **kwargs):
|
||||
new_srcs: list[UOp] = [UOp.const(x) if isinstance(x, int) else x for x in srcs if x is not None]
|
||||
if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].val]
|
||||
@@ -814,7 +822,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
unique_num = itertools.count(0)
|
||||
|
||||
def getaddr(self, device=None) -> UOp:
|
||||
if self.without_after.op not in {Ops.BUFFER, Ops.SHRINK, Ops.BITCAST, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
|
||||
if self.without_after.op not in {Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
|
||||
return UOp(Ops.GETADDR, src=(self,), arg=device or to_tuple(self.device)[0])
|
||||
@staticmethod
|
||||
def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
|
||||
@@ -893,7 +901,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
while len(s.src) and s.op not in {Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.MSTACK}: s = s.src[0]
|
||||
return s
|
||||
|
||||
def contiguous_view(self) -> tuple[UOp, int]|None:
|
||||
def contiguous_view_offset(self) -> int|None:
|
||||
"""If movement ops on a BUFFER collapse to a contiguous range, return `offset` in elements. Otherwise None."""
|
||||
from tinygrad.schedule.rangeify import pm_mops
|
||||
from tinygrad.uop.symbolic import symbolic
|
||||
|
||||
@@ -906,17 +915,14 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
|
||||
idx = self.flatten().index(UOp.range(self.numel(), 0))
|
||||
out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset")
|
||||
if out.op is not Ops.INDEX or not (b:=out.src[0]).tag or (c:=out.src[1]).op is not Ops.CONST or not isinstance(c.val, int): return None
|
||||
return b.rtag(None), c.val
|
||||
|
||||
def contiguous_view_offset(self) -> int|None: return None if (view := self.contiguous_view()) is None else view[1]
|
||||
return out.val if out.op is Ops.CONST and isinstance(out.val, int) else None
|
||||
|
||||
def has_buffer_identity(self, after_ok=False):
|
||||
"""Check if this UOp has a concrete buffer identity in the graph (RESHAPE/UNSHARD -> BUFFER chain)."""
|
||||
# TODO: this is confusing because UOp.variable('v', 0, 1, dtypes.weakfloat) is True for jit to work, but it doesn't have a buffer
|
||||
if self.op in {Ops.RESHAPE, Ops.UNSHARD, Ops.MSELECT}: return self.src[0].has_buffer_identity(after_ok)
|
||||
if after_ok and self.op == Ops.AFTER: return self.src[0].has_buffer_identity(after_ok)
|
||||
return self.op in {Ops.BUFFER, Ops.PARAM}
|
||||
return self.op in {Ops.BUFFER, Ops.SLICE, Ops.PARAM}
|
||||
|
||||
def _base_buffer_is_realized(self) -> bool:
|
||||
"""Walk through AFTER chain to find if the underlying buffer is realized (has allocated memory)."""
|
||||
@@ -926,19 +932,30 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
|
||||
@property
|
||||
def buffer(self) -> Buffer|MultiBuffer:
|
||||
if self.op in {Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
|
||||
if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
|
||||
# this buffer can process disk tensors and simple movement ops
|
||||
if self is not self.base or self.op is Ops.BITCAST:
|
||||
if self is not self.base:
|
||||
buf = self.base.buffer
|
||||
assert isinstance(buf, Buffer), "must be a Buffer for movement ops"
|
||||
offset = self.contiguous_view_offset()
|
||||
if offset is None: raise RuntimeError(f"non-contiguous view is not supported for {buf.device} buffer")
|
||||
return buf.view(prod(self.max_shape), self.dtype, offset*self.dtype.itemsize)
|
||||
if self.op is Ops.BITCAST:
|
||||
buf = self.src[0].buffer
|
||||
assert isinstance(buf, Buffer), "must be a Buffer for BITCAST"
|
||||
return buf.view(prod(self.max_shape), self.dtype, 0)
|
||||
if self.op is Ops.SLICE:
|
||||
if (cret:=buffers.get(self)) is not None: return cret
|
||||
if (cv := self.contiguous_view()) is None: raise RuntimeError(f"non-contiguous view is not supported for {self.device} buffer")
|
||||
buf, offset = (b:=cv[0]).base.buffer, cv[1]
|
||||
buf = self.src[0].buffer
|
||||
offset = self.src[1].val
|
||||
if isinstance(buf, MultiBuffer):
|
||||
mbuf = MultiBuffer.__new__(MultiBuffer)
|
||||
mbuf.bufs = [x.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize) for x in buf.bufs]
|
||||
mbuf.bufs = [b.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize) for b in buf.bufs]
|
||||
buffers[self] = mbuf
|
||||
return mbuf
|
||||
buffers[self] = buf.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize)
|
||||
return buffers[self]
|
||||
assert isinstance(buf, Buffer), "must be a Buffer for SLICE"
|
||||
buffers[self] = bv = buf.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize)
|
||||
return bv
|
||||
if self.op is Ops.MSELECT:
|
||||
ret = self.src[0].buffer
|
||||
assert isinstance(ret, MultiBuffer)
|
||||
@@ -1164,7 +1181,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)
|
||||
|
||||
# opaque bodies stay as Ops.CALL; value-producing bodies become Ops.FUNCTION (wrapped in TUPLE)
|
||||
_OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.CUSTOM_FUNCTION}
|
||||
_OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.SLICE, Ops.CUSTOM_FUNCTION}
|
||||
def call(self, *srcs:UOp, ret_dtype:DType|None=None, grad_fxn:Callable|None=None,
|
||||
name:str|None=None, precompile:bool=False, precompile_backward:bool=False, aux:Any=None) -> UOp:
|
||||
if ret_dtype is not None: return UOp(Ops.CALL, ret_dtype, src=(self,)+srcs)
|
||||
@@ -1758,14 +1775,10 @@ pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)])
|
||||
|
||||
# ctx is source UOp for which we are finding a contiguous view for. used in contiguous_view_offset
|
||||
pm_contiguous_view_offset = PatternMatcher([
|
||||
# normalize to 1d bitcasts
|
||||
(UPat(Ops.BITCAST, name="b"), lambda b: b.src[0].flatten().bitcast(b.dtype).reshape(b.shape) if len(b.shape) != 1 else None),
|
||||
(UPat(Ops.BITCAST, name="b").index(UPat.cvar("c")), lambda ctx, b, c:
|
||||
b.src[0].flatten().index(UOp.range(ctx.numel() * (osz:=b.element_size())//(isz:=b.src[0].element_size()), 0) + (c * osz//isz)) if b.tag else None),
|
||||
(UPat(Ops.INDEX, src=(UPat.var("b"),)), lambda b: b.rtag().index(0)),
|
||||
(UPat(Ops.INDEX, src=(UPat.var("b"), UPat(Ops.RANGE))), lambda b: b.rtag().index(0)),
|
||||
(UPat(Ops.INDEX, src=(UPat.var("b"), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda ctx, b, c: b.rtag().index(c)),
|
||||
(UPat(Ops.INDEX, src=(UPat.var("b"), UPat.cvar('c'))), lambda ctx, b, c: b.rtag().index(c) if resolve(ctx.numel() == 1, False) else None),
|
||||
(UPat(Ops.INDEX, src=(UPat(),)), lambda: UOp.const(0)),
|
||||
(UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE))), lambda: UOp.const(0)),
|
||||
(UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda c: c),
|
||||
(UPat(Ops.INDEX, src=(UPat(), UPat.cvar('c'))), lambda ctx, c: c if resolve(ctx.numel() == 1, False) else None),
|
||||
])
|
||||
|
||||
# *** what was symbolic.py ***
|
||||
|
||||
@@ -233,6 +233,13 @@ spec_hcq = PatternMatcher([
|
||||
spec_full = PatternMatcher([
|
||||
(UPat(Ops.REWRITE_ERROR, dtypes.void, name="x"), lambda x: isinstance(x.arg, str)),
|
||||
|
||||
# SLICE on BUFFER is allowed if BUFFER is
|
||||
(UPat(Ops.SLICE, src=(UPat(GroupOp.Movement.union({Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.AFTER})),
|
||||
UPat(Ops.CONST, dtype=dtypes.weakint)), allow_any_len=True, name="bv"),
|
||||
lambda bv: isinstance(bv.arg, int)),
|
||||
|
||||
(UPat(Ops.CALL, dtypes.void, src=(UPat((Ops.SLICE,)),), allow_any_len=True), lambda: True),
|
||||
|
||||
# codegen may end ranges after gpudims has replaced RANGE with SPECIAL.
|
||||
(UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True), lambda: True),
|
||||
|
||||
@@ -268,7 +275,7 @@ spec_kernel_graph = PatternMatcher([
|
||||
(UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)),
|
||||
(UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
|
||||
# all calls are on various sinks
|
||||
(UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True),
|
||||
(UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM)),), allow_any_len=True), lambda: True),
|
||||
# after on PARAM or AFTER
|
||||
(UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST, Ops.RESHAPE})),),
|
||||
allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)),
|
||||
@@ -279,11 +286,10 @@ spec_kernel_graph = PatternMatcher([
|
||||
# late imports to avoid circular import
|
||||
from tinygrad.codegen.opt import Opt, OptOps
|
||||
from tinygrad.schedule.rangeify import BufferizeOpts
|
||||
from tinygrad.renderer import Estimates
|
||||
glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Metadata": Metadata,
|
||||
"UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
|
||||
"Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace, "panic": panic,
|
||||
"ConstFloat": ConstFloat, "ParamArg": ParamArg, "Estimates": Estimates}
|
||||
"ConstFloat": ConstFloat, "ParamArg": ParamArg}
|
||||
def eval_pyrender(code:str) -> UOp:
|
||||
lcls:dict[str, Any] = {}
|
||||
exec(code, glbls, lcls)
|
||||
|
||||
@@ -249,8 +249,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
|
||||
lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
|
||||
(UPat(Ops.RANGE, src=(UPat(Ops.CONST,)), name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
|
||||
# max folding
|
||||
((UPat.cvar("a") < UPat.var("b")).where(UPat.var("b"), UPat.cvar("c")), lambda a,b,c: UOp.maximum(a,b) if a.val == c.val else None),
|
||||
((UPat.var("a") < UPat.cvar("b")).where(UPat.cvar("c"), UPat.var("a")), lambda a,b,c: UOp.maximum(a,b) if b.val == c.val else None),
|
||||
(UPat.maximum(UPat.var("x"), UPat.var("y")), lambda x,y: x if x.vmin >= y.vmax else y if x.vmax <= y.vmin else None),
|
||||
# TODO: why does this rule break beautiful_mnist?
|
||||
#((UPat.var("x")+UPat.var("z")).maximum(UPat.var("y")+UPat.var("z")), lambda x,y,z: x.maximum(y) + z),
|
||||
@@ -295,8 +293,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
|
||||
(UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
|
||||
tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.LINEAR, Ops.STAGE}
|
||||
else y.src for y in x.src[1:]]))))),
|
||||
# after/end with 1 src is just src[0]
|
||||
(UPat((Ops.AFTER, Ops.END), src=(UPat.var("s"),)), lambda s: s),
|
||||
# after with 1 src is just src[0]
|
||||
(UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
|
||||
])+div_and_mod_symbolic
|
||||
|
||||
# ******** we take a small aside to "simplify_valid" to rewrite valids ********
|
||||
|
||||
@@ -50,7 +50,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
|
||||
Ops.INDEX: "#CEF9B7", Ops.STACK: "#D8F9E4",
|
||||
Ops.WMMA: "#efefc0", Ops.UNSHARD: "#f6ccff", Ops.INS: "#eec4ff",
|
||||
**{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
|
||||
Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
|
||||
Ops.SLICE: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
|
||||
Ops.CALL: "#00B7C8", Ops.FUNCTION: "#C07788", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040",
|
||||
Ops.LINEAR: "#7DF4FF",
|
||||
Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
|
||||
|
||||
Reference in New Issue
Block a user