sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

Merge branch 'master' into remove_assign

Browse Source
This commit is contained in:
George Hotz
2025-11-14 09:16:27 -08:00
committed by GitHub
parent c793a08fbf 7c110e1a57
commit 0dd7c13aa8
110 changed files with 133665 additions and 208554 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/actions/setup-tinygrad/action.yml
+5 -5
View File
@@ -61,7 +61,7 @@ runs:
uses: actions/cache@v4
with:
path: ${{ github.workspace }}/.venv
key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ hashFiles('**/pyproject.toml') }}-${{ env.PYTHON_CACHE_VERSION }}
key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ hashFiles('**/pyproject.toml') }}-${{ env.CACHE_VERSION }}
# **** Caching downloads ****
@@ -70,13 +70,13 @@ runs:
uses: actions/cache@v4
with:
path: ~/.cache/tinygrad/downloads/
key: downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}
key: downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}
- name: Cache downloads (macOS)
if: inputs.key != '' && runner.os == 'macOS'
uses: actions/cache@v4
with:
path: ~/Library/Caches/tinygrad/downloads/
key: osx-downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}
key: osx-downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}
# **** Python deps ****
@@ -187,7 +187,7 @@ runs:
uses: actions/cache@v4
with:
path: /var/cache/apt/archives/
key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.APT_CACHE_VERSION }}
key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
- name: Run apt Update + Install
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
@@ -247,7 +247,7 @@ runs:
cache-name: cache-gpuocelot-build-1
with:
path: ${{ github.workspace }}/gpuocelot/ocelot
key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
- name: Clone/compile gpuocelot
if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
shell: bash
.github/workflows/autogen.yml
+94 -49
View File
@@ -1,10 +1,7 @@
name: Autogen
env:
# increment this when downloads substantially change to avoid the internet
DOWNLOAD_CACHE_VERSION: '12'
PYTHON_CACHE_VERSION: '4'
APT_CACHE_VERSION: '1'
BUILD_CACHE_VERSION: '1'
CACHE_VERSION: '13'
CAPTURE_PROCESS_REPLAY: 1
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PYTHONPATH: ${{ github.workspace }}
@@ -14,15 +11,15 @@ on:
branches:
- master
pull_request:
paths:
paths:
- 'tinygrad/runtime/autogen/**/*'
workflow_dispatch:
paths:
paths:
- 'tinygrad/runtime/autogen/**/*'
jobs:
autogen:
name: Autogen
name: In-tree Autogen
runs-on: ubuntu-24.04
timeout-minutes: 15
steps:
@@ -34,66 +31,107 @@ jobs:
opencl: 'true'
amd: 'true'
cuda: 'true'
webgpu: 'true'
llvm: 'true'
pydeps: 'pyyaml mako'
webgpu: 'true'
mesa: 'true'
pydeps: 'clang>=20 pyyaml mako'
- name: Install autogen support packages
run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev llvm-20-dev
run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
- name: Verify OpenCL autogen
run: |
cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
./autogen_stubs.sh opencl
mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
python3 -c "from tinygrad.runtime.autogen import opencl"
diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
- name: Verify CUDA autogen
run: |
cp tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
cp tinygrad/runtime/autogen/nv_gpu.py /tmp/nv_gpu.py.bak
./autogen_stubs.sh cuda
./autogen_stubs.sh nv
mv tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
mv tinygrad/runtime/autogen/nvrtc.py /tmp/nvrtc.py.bak
mv tinygrad/runtime/autogen/nvjitlink.py /tmp/nvjitlink.py.bak
mv tinygrad/runtime/autogen/nv_570.py /tmp/nv_570.py.bak
mv tinygrad/runtime/autogen/nv.py /tmp/nv.py.bak
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
diff /tmp/nv_gpu.py.bak tinygrad/runtime/autogen/nv_gpu.py
diff /tmp/nvrtc.py.bak tinygrad/runtime/autogen/nvrtc.py
diff /tmp/nvjitlink.py.bak tinygrad/runtime/autogen/nvjitlink.py
diff /tmp/nv_570.py.bak tinygrad/runtime/autogen/nv_570.py
diff /tmp/nv.py.bak tinygrad/runtime/autogen/nv.py
- name: Verify AMD autogen
run: |
cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
./autogen_stubs.sh hsa
./autogen_stubs.sh kfd
./autogen_stubs.sh comgr
./autogen_stubs.sh amd
./autogen_stubs.sh sqtt
diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
mv tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
mv tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
mv tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
mv tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
mv tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
mv tinygrad/runtime/autogen/rocprof.py /tmp/rocprof.py.bak
mv tinygrad/runtime/autogen/am/am.py /tmp/am_am.py.bak
mv tinygrad/runtime/autogen/am/pm4_soc15.py /tmp/am_pm4_soc15.py.bak
mv tinygrad/runtime/autogen/am/pm4_nv.py /tmp/am_pm4_nv.py.bak
mv tinygrad/runtime/autogen/am/sdma_4_0_0.py /tmp/am_sdma_4_0_0.py.bak
mv tinygrad/runtime/autogen/am/sdma_5_0_0.py /tmp/am_sdma_5_0_0.py.bak
mv tinygrad/runtime/autogen/am/sdma_6_0_0.py /tmp/am_sdma_6_0_0.py.bak
mv tinygrad/runtime/autogen/am/smu_v13_0_0.py /tmp/am_smu_v13_0_0.py.bak
mv tinygrad/runtime/autogen/am/smu_v14_0_2.py /tmp/am_smu_v14_0_2.py.bak
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof; from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
diff /tmp/hip.py.bak tinygrad/runtime/autogen/hip.py
diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
diff /tmp/rocprof.py.bak tinygrad/runtime/autogen/rocprof.py
diff /tmp/am_am.py.bak tinygrad/runtime/autogen/am/am.py
diff /tmp/am_pm4_soc15.py.bak tinygrad/runtime/autogen/am/pm4_soc15.py
diff /tmp/am_pm4_nv.py.bak tinygrad/runtime/autogen/am/pm4_nv.py
diff /tmp/am_sdma_4_0_0.py.bak tinygrad/runtime/autogen/am/sdma_4_0_0.py
diff /tmp/am_sdma_5_0_0.py.bak tinygrad/runtime/autogen/am/sdma_5_0_0.py
diff /tmp/am_sdma_6_0_0.py.bak tinygrad/runtime/autogen/am/sdma_6_0_0.py
diff /tmp/am_smu_v13_0_0.py.bak tinygrad/runtime/autogen/am/smu_v13_0_0.py
diff /tmp/am_smu_v14_0_2.py.bak tinygrad/runtime/autogen/am/smu_v14_0_2.py
- name: Verify Linux autogen
run: |
cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
./autogen_stubs.sh io_uring
./autogen_stubs.sh ib
mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
mv tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
mv tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
mv tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
mv tinygrad/runtime/autogen/pci.py /tmp/pci.py.bak
mv tinygrad/runtime/autogen/vfio.py /tmp/vfio.py.bak
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
- name: Verify WebGPU autogen
run: |
cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
./autogen_stubs.sh webgpu
diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
diff /tmp/pci.py.bak tinygrad/runtime/autogen/pci.py
diff /tmp/vfio.py.bak tinygrad/runtime/autogen/vfio.py
- name: Verify LLVM autogen
run: |
cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
./autogen_stubs.sh llvm
mv tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
python3 -c "from tinygrad.runtime.autogen import llvm"
diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
- name: Verify WebGPU autogen
run: |
mv tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
python3 -c "from tinygrad.runtime.autogen import webgpu"
diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
- name: Verify Qualcomm autogen
run: |
mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
mv tinygrad/runtime/autogen/adreno.py /tmp/adreno.py.bak
mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
python3 -c "from tinygrad.runtime.autogen import kgsl, adreno, qcom_dsp"
diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
diff /tmp/adreno.py.bak tinygrad/runtime/autogen/adreno.py
diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
- name: Verify libusb autogen
run: |
mv tinygrad/runtime/autogen/libusb.py /tmp/libusb.py.bak
python3 -c "from tinygrad.runtime.autogen import libusb"
diff /tmp/libusb.py.bak tinygrad/runtime/autogen/libusb.py
- name: Verify mesa autogen
run: |
cp tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
./autogen_stubs.sh mesa
mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
python3 -c "from tinygrad.runtime.autogen import mesa"
diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
autogen-ng:
name: In-tree Autogen
autogen-comgr-3:
name: In-tree Autogen (comgr 3)
runs-on: ubuntu-24.04
timeout-minutes: 15
steps:
@@ -104,9 +142,16 @@ jobs:
with:
pydeps: 'clang>=20'
- name: Install autogen support packages
run: sudo apt-get install -y --no-install-recommends libclang-20-dev
- name: Verify Linux autogen
run: |
mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
python3 -c "from tinygrad.runtime.autogen import libc"
diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.4 $(lsb_release -cs) main
EOF
echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
sudo apt -qq update || true
sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
- name: Verify comgr (3) autogen
run: |
mv tinygrad/runtime/autogen/comgr_3.py /tmp/comgr_3.py.bak
python3 -c "from tinygrad.runtime.autogen import comgr_3"
diff /tmp/comgr_3.py.bak tinygrad/runtime/autogen/comgr_3.py
.github/workflows/benchmark.yml
+28 -26
View File
@@ -54,7 +54,7 @@ jobs:
- name: Print macOS version
run: sw_vers
- name: Run Stable Diffusion
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
- name: Run Stable Diffusion without fp16
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
- name: Run Stable Diffusion v2
@@ -199,7 +199,7 @@ jobs:
- name: Test speed vs torch
run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
- name: Test speed vs theoretical
run: NV=1 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
run: NV=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
- name: Test benchmark allreduce
run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
- name: Test tensor cores
@@ -320,19 +320,20 @@ jobs:
# run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
- name: Train MNIST
run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
# TODO: too slow
- name: Run 10 CIFAR training steps
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
- name: Run 10 CIFAR training steps w HALF
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
- name: Run 10 CIFAR training steps w BF16
run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=1300 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
# - name: Run 10 CIFAR training steps w HALF
# run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
# - name: Run 10 CIFAR training steps w BF16
# run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
# TODO: too slow
# - name: Run 10 CIFAR training steps w winograd
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
- name: Run full CIFAR training steps w 6 GPUS
run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
# - name: Run full CIFAR training w 1 GPU
# run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
# - name: Run full CIFAR training steps w 6 GPUS
# run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
- name: Run MLPerf resnet eval on training data
run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
#- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
@@ -409,7 +410,7 @@ jobs:
# python3 -c "import torch; print(torch.__version__)"
# LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
- name: Test speed vs theoretical
run: AMD=1 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
- name: Test tensor cores
run: |
AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
@@ -524,17 +525,18 @@ jobs:
run: test/external/process_replay/reset.py
- name: Train MNIST
run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
# TODO: too slow
- name: Run 10 CIFAR training steps
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
- name: Run 10 CIFAR training steps w HALF
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=390 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=2000 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
# - name: Run 10 CIFAR training steps w HALF
# run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=390 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
# - name: Run 10 CIFAR training steps w BF16
# run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
# TODO: too slow
# - name: Run 10 CIFAR training steps w winograd
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
# - name: Run full CIFAR training w 1 GPU
# run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
#- name: Run full CIFAR training steps w 6 GPUS
# run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
#- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
@@ -632,17 +634,15 @@ jobs:
- name: openpilot compile3 0.10.0 driving_policy
run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
- name: openpilot compile3 0.10.0 dmonitoring
run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
- name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.10.1 driving_vision
# TODO: ASSERT_MIN_STEP_TIME=17
run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=21 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.10.1 driving_policy
run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
- name: openpilot compile3 0.10.1 dmonitoring
# TODO: ASSERT_MIN_STEP_TIME=10
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
- name: benchmark MobileNetV2 on DSP
run: |
# generate quantized weights
@@ -706,8 +706,9 @@ jobs:
run: |
AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
# TODO: too slow
# - name: Run full CIFAR training w 1 GPU
# run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
# TODO: enable
# - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
# run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
@@ -769,8 +770,9 @@ jobs:
NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
- name: Test LLAMA-3
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
# TODO: too slow
# - name: Run full CIFAR training w 1 GPU
# run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
#- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
# run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
- name: Run 10 MLPerf Bert training steps (1 gpu)
.github/workflows/benchmark_search.yml
+2 -2
View File
@@ -22,13 +22,13 @@ jobs:
- name: Run SDXL with new search
# TODO: GCVM_L2_PROTECTION_FAULT_STATUS with llvm19
run: |
BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 python examples/sdxl.py --noshow --timing --seed 0
BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CCACHE=0 python examples/sdxl.py --noshow --timing --seed 0
- name: Run SDXL with cached search
run: |
BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
- name: Run winograd cifar with new search
run: |
BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 CCACHE=0 BS=1024 STEPS=500 python examples/hlb_cifar10.py
- name: Run winograd cifar with cached search
run: |
BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py
.github/workflows/test.yml
+6 -8
View File
@@ -1,10 +1,7 @@
name: Unit Tests
env:
# increment this when downloads substantially change to avoid the internet
DOWNLOAD_CACHE_VERSION: '12'
PYTHON_CACHE_VERSION: '4'
APT_CACHE_VERSION: '1'
BUILD_CACHE_VERSION: '1'
CACHE_VERSION: '13'
CAPTURE_PROCESS_REPLAY: 1
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PYTHONPATH: ${{ github.workspace }}
@@ -233,7 +230,7 @@ jobs:
python-version: '3.11'
deps: linting
- name: Lint bad-indentation and trailing-whitespace with pylint
run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string=' ' --recursive=y . --ignore-paths='tinygrad/runtime/autogen'
run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string=' ' --recursive=y .
- name: Lint with ruff
run: |
pip3 install --upgrade --force-reinstall ruff==0.11.0
@@ -344,10 +341,11 @@ jobs:
key: gpu-image
deps: testing_minimal
opencl: 'true'
- name: Test CL IMAGE=2 ops + training
- name: Test CL IMAGE=2 ops
run: |
CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
# TODO: training is broken
# CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -392,7 +390,7 @@ jobs:
llvm: 'true'
- name: Test openpilot model kernel count and gate usage
run: |
ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1452 ALLOWED_GATED_READ_IMAGE=122 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
- name: Test openpilot CL compile fp16
run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
- name: Test openpilot CL compile fp32 (test correctness)
README.md
+27 -6
View File
@@ -21,17 +21,38 @@ tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) an
---
Despite tinygrad's size, it is a fully featured deep learning framework.
tinygrad is an end-to-end deep learning stack:
Due to its extreme simplicity, it is the easiest framework to add new accelerators to, with support for both inference and training. If XLA is CISC, tinygrad is RISC.
- **Tensor library** with autograd
- **IR and compiler** that fuse and lower kernels
- **JIT + graph execution**
- **nn / optim / datasets** for real training
tinygrad is now beta software, we [raised some money](https://geohot.github.io/blog/jekyll/update/2023/05/24/the-tiny-corp-raised-5M.html) to make it good. Someday, we will tape out chips.
It’s inspired by PyTorch (ergonomics), JAX (functional transforms and IR-based AD), and TVM (scheduling and codegen), but stays intentionally tiny and hackable.
## Features
---
### LLaMA and Stable Diffusion
## How tinygrad compares
tinygrad can run [LLaMA](/docs/showcase.md#llama) and [Stable Diffusion](/docs/showcase.md#stable-diffusion)!
**PyTorch**
- ✅ Similar: eager `Tensor` API, autograd, `optim`, basic datasets and layers.
- ✅ You can write familiar training loops.
- 🔁 Unlike PyTorch, the entire compiler and IR are visible and hackable.
**JAX**
- ✅ IR-based autodiff over primitives (like JAXPR + XLA).
- ✅ Function-level JIT (`TinyJit`) that captures and replays kernels.
- 🔁 Fewer functional transforms (no full `vmap`/`pmap` yet), but far easier to read.
**TVM**
- ✅ Multiple lowering passes, scheduling, and BEAM search over kernels.
- ✅ Device “graphs” for batched execution.
- 🔁 tinygrad also ships the **front-end framework** (tensors, nn, optim), not just the compiler.
---
### Laziness
autogen_stubs.sh
-550
View File
@@ -1,550 +0,0 @@
#!/bin/bash -e
# setup instructions for clang2py
if [[ ! $(clang2py -V) ]]; then
pushd .
cd /tmp
sudo apt-get install -y --no-install-recommends clang
pip install --upgrade pip setuptools
pip install clang==14.0.6
git clone https://github.com/nimlgen/ctypeslib.git
cd ctypeslib
pip install .
clang2py -V
popd
fi
BASE=tinygrad/runtime/autogen/
fixup() {
sed -i '1s/^/# mypy: ignore-errors\n/' $1
sed -i 's/ *$//' $1
grep FIXME_STUB $1 || true
}
patch_dlopen() {
path=$1; shift
name=$1; shift
cat <<EOF | sed -i "/import ctypes.*/r /dev/stdin" $path
PATHS_TO_TRY = [
$(for p in "$@"; do echo " $p,"; done)
]
def _try_dlopen_$name():
library = ctypes.util.find_library("$name")
if library:
try: return ctypes.CDLL(library)
except OSError: pass
for candidate in PATHS_TO_TRY:
try: return ctypes.CDLL(candidate)
except OSError: pass
return None
EOF
}
generate_opencl() {
clang2py /usr/include/CL/cl.h -o $BASE/opencl.py -l /usr/lib/x86_64-linux-gnu/libOpenCL.so.1 -k cdefstum
fixup $BASE/opencl.py
# hot patches
sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/opencl.py
sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libOpenCL.so.1')\ctypes.CDLL(ctypes.util.find_library('OpenCL'))\g" $BASE/opencl.py
python3 -c "import tinygrad.runtime.autogen.opencl"
}
generate_hip() {
clang2py /opt/rocm/include/hip/hip_ext.h /opt/rocm/include/hip/hiprtc.h \
/opt/rocm/include/hip/hip_runtime_api.h /opt/rocm/include/hip/driver_types.h \
--clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/hip.py -l /opt/rocm/lib/libamdhip64.so
echo "hipDeviceProp_t = hipDeviceProp_tR0600" >> $BASE/hip.py
echo "hipGetDeviceProperties = hipGetDevicePropertiesR0600" >> $BASE/hip.py
fixup $BASE/hip.py
# we can trust HIP is always at /opt/rocm/lib
#sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/hip.py
#sed -i "s\ctypes.CDLL('/opt/rocm/lib/libhiprtc.so')\ctypes.CDLL(ctypes.util.find_library('hiprtc'))\g" $BASE/hip.py
#sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamdhip64.so')\ctypes.CDLL(ctypes.util.find_library('amdhip64'))\g" $BASE/hip.py
sed -i "s\import ctypes\import ctypes, os\g" $BASE/hip.py
sed -i "s\'/opt/rocm/\os.getenv('ROCM_PATH', '/opt/rocm/')+'/\g" $BASE/hip.py
python3 -c "import tinygrad.runtime.autogen.hip"
}
generate_comgr() {
clang2py /opt/rocm/include/amd_comgr/amd_comgr.h \
--clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/comgr.py -l /opt/rocm/lib/libamd_comgr.so
fixup $BASE/comgr.py
sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/comgr.py
patch_dlopen $BASE/comgr.py amd_comgr "'/opt/rocm/lib/libamd_comgr.so'" "os.getenv('ROCM_PATH', '')+'/lib/libamd_comgr.so'" "'/usr/local/lib/libamd_comgr.dylib'" "'/opt/homebrew/lib/libamd_comgr.dylib'"
sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamd_comgr.so')\_try_dlopen_amd_comgr()\g" $BASE/comgr.py
python3 -c "import tinygrad.runtime.autogen.comgr"
}
generate_kfd() {
clang2py /usr/include/linux/kfd_ioctl.h -o $BASE/kfd.py -k cdefstum
fixup $BASE/kfd.py
sed -i "s/import ctypes/import ctypes, os/g" $BASE/kfd.py
sed -i "s/import fcntl, functools/import functools/g" $BASE/kfd.py
sed -i "/import functools/a from tinygrad.runtime.support.hcq import FileIOInterface" $BASE/kfd.py
sed -i "s/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd, \*\*kwargs):/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd:FileIOInterface, \*\*kwargs):/g" $BASE/kfd.py
sed -i "s/fcntl.ioctl(__fd, (__idir<<30)/__fd.ioctl((__idir<<30)/g" $BASE/kfd.py
sed -i "s/!!/not not /g" $BASE/kfd.py
python3 -c "import tinygrad.runtime.autogen.kfd"
}
generate_cuda() {
clang2py /usr/include/cuda.h --clang-args="-D__CUDA_API_VERSION_INTERNAL" -o $BASE/cuda.py -l /usr/lib/x86_64-linux-gnu/libcuda.so
sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/cuda.py
sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libcuda.so')\ctypes.CDLL(ctypes.util.find_library('cuda'))\g" $BASE/cuda.py
fixup $BASE/cuda.py
python3 -c "import tinygrad.runtime.autogen.cuda"
}
generate_nvrtc() {
clang2py /usr/local/cuda/include/nvrtc.h /usr/local/cuda/include/nvJitLink.h -o $BASE/nvrtc.py -l /usr/local/cuda/lib64/libnvrtc.so -l /usr/local/cuda/lib64/libnvJitLink.so
sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/nvrtc.py
sed -i "s\ctypes.CDLL('/usr/local/cuda/lib64/libnvrtc.so')\ctypes.CDLL(ctypes.util.find_library('nvrtc'))\g" $BASE/nvrtc.py
sed -i "s\ctypes.CDLL('/usr/local/cuda/lib64/libnvJitLink.so')\ctypes.CDLL(ctypes.util.find_library('nvJitLink'))\g" $BASE/nvrtc.py
fixup $BASE/nvrtc.py
python3 -c "import tinygrad.runtime.autogen.nvrtc"
}
generate_nv() {
NVKERN_COMMIT_HASH=81fe4fb417c8ac3b9bdcc1d56827d116743892a5
NVKERN_SRC=/tmp/open-gpu-kernel-modules-$NVKERN_COMMIT_HASH
if [ ! -d "$NVKERN_SRC" ]; then
git clone https://github.com/NVIDIA/open-gpu-kernel-modules $NVKERN_SRC
pushd .
cd $NVKERN_SRC
git reset --hard $NVKERN_COMMIT_HASH
popd
fi
clang2py -k cdefstum \
extra/nv_gpu_driver/clc6c0qmd.h \
extra/nv_gpu_driver/clcec0qmd.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0000.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0080.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080_notification.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc56f.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc86f.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc96f.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc761.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl83de.h \
$NVKERN_SRC/src/nvidia/generated/g_allclasses.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc6c0.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/class/clcdc0.h \
$NVKERN_SRC/kernel-open/nvidia-uvm/clc6b5.h \
$NVKERN_SRC/kernel-open/nvidia-uvm/clc9b5.h \
$NVKERN_SRC/kernel-open/nvidia-uvm/uvm_ioctl.h \
$NVKERN_SRC/kernel-open/nvidia-uvm/uvm_linux_ioctl.h \
$NVKERN_SRC/kernel-open/nvidia-uvm/hwref/ampere/ga100/dev_fault.h \
$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv_escape.h \
$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl.h \
$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl-numbers.h \
$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl-numa.h \
$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-unix-nvos-params-wrappers.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/alloc/alloc_channel.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/nvos.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl0000/*.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl0080/*.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl2080/*.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl83de/*.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlc36f.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlcb33.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrla06c.h \
$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl90f1.h \
--clang-args="-include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
-o $BASE/nv_gpu.py
fixup $BASE/nv_gpu.py
sed -i "s\(0000000001)\1\g" $BASE/nv_gpu.py
sed -i "s\import ctypes\import ctypes, os\g" $BASE/nv_gpu.py
sed -i 's/#\?\s\([A-Za-z0-9_]\+\) = MW ( \([0-9]\+\) : \([0-9]\+\) )/\1 = (\2 , \3)/' $BASE/nv_gpu.py # NVC6C0_QMDV03_00 processing
sed -i 's/#\sdef NVC6C0_QMD\([A-Za-z0-9_()]\+\):/def NVC6C0_QMD\1:/' $BASE/nv_gpu.py
sed -i 's/#\sdef NVCEC0_QMD\([A-Za-z0-9_()]\+\):/def NVCEC0_QMD\1:/' $BASE/nv_gpu.py
sed -E -i -n '/^def (NVCEC0_QMDV05_00_RELEASE)(_ENABLE)\(i\):/{p;s//\1'"0"'\2=\1\2(0)\n\1'"1"'\2=\1\2(1)/;H;b};p;${x;s/^\n//;p}' "$BASE/nv_gpu.py"
sed -i 's/#\s*return MW(\([0-9i()*+]\+\):\([0-9i()*+]\+\))/ return (\1 , \2)/' $BASE/nv_gpu.py
sed -i 's/#\?\s*\(.*\)\s*=\s*\(NV\)\?BIT\(32\)\?\s*(\s*\([0-9]\+\)\s*)/\1 = (1 << \4)/' $BASE/nv_gpu.py # name = BIT(x) -> name = (1 << x)
sed -i "s/UVM_\([A-Za-z0-9_]\+\) = \['i', '(', '\([0-9]\+\)', ')'\]/UVM_\1 = \2/" $BASE/nv_gpu.py # UVM_name = ['i', '(', '<num>', ')'] -> UVM_name = <num>
# Parse status codes
sed -n '1i\
nv_status_codes = {}
/^NV_STATUS_CODE/ { s/^NV_STATUS_CODE(\([^,]*\), *\([^,]*\), *"\([^"]*\)") *.*$/\1 = \2\nnv_status_codes[\1] = "\3"/; p }' $NVKERN_SRC/src/common/sdk/nvidia/inc/nvstatuscodes.h >> $BASE/nv_gpu.py
python3 -c "import tinygrad.runtime.autogen.nv_gpu"
clang2py -k cdefstum \
$NVKERN_SRC/src/nvidia/inc/kernel/gpu/fsp/kern_fsp_cot_payload.h \
$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gspifpub.h \
$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_wpr_meta.h \
$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_sr_meta.h \
$NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
$NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
$NVKERN_SRC/src/common/uproc/os/common/include/libos_init_args.h \
$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/rmRiscvUcode.h \
$NVKERN_SRC/src/common/shared/msgq/inc/msgq/msgq_priv.h \
$NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_headers.h \
$NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h \
$NVKERN_SRC/src/nvidia/generated/g_rpc-structures.h \
$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/fsp/fsp_nvdm_format.h \
extra/nv_gpu_driver/g_rpc-message-header.h \
extra/nv_gpu_driver/gsp_static_config.h \
extra/nv_gpu_driver/vbios.h \
extra/nv_gpu_driver/pci_exp_table.h \
--clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
-o $BASE/nv/nv.py
fixup $BASE/nv/nv.py
python3 -c "import tinygrad.runtime.autogen.nv.nv"
}
generate_amd() {
# clang2py broken when pass -x c++ to prev headers
clang2py -k cdefstum \
extra/hip_gpu_driver/sdma_registers.h \
extra/hip_gpu_driver/nvd.h \
extra/hip_gpu_driver/gc_11_0_0_offset.h \
extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
--clang-args="-I/opt/rocm/include -x c++" \
-o $BASE/amd_gpu.py
fixup $BASE/amd_gpu.py
sed -i "s\import ctypes\import ctypes, os\g" $BASE/amd_gpu.py
python3 -c "import tinygrad.runtime.autogen.amd_gpu"
}
generate_hsa() {
clang2py \
/opt/rocm/include/hsa/hsa.h \
/opt/rocm/include/hsa/hsa_ext_amd.h \
/opt/rocm/include/hsa/amd_hsa_signal.h \
/opt/rocm/include/hsa/amd_hsa_queue.h \
/opt/rocm/include/hsa/amd_hsa_kernel_code.h \
/opt/rocm/include/hsa/hsa_ext_finalize.h /opt/rocm/include/hsa/hsa_ext_image.h \
/opt/rocm/include/hsa/hsa_ven_amd_aqlprofile.h \
--clang-args="-I/opt/rocm/include" \
-o $BASE/hsa.py -l /opt/rocm/lib/libhsa-runtime64.so
fixup $BASE/hsa.py
sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/hsa.py
sed -i "s\ctypes.CDLL('/opt/rocm/lib/libhsa-runtime64.so')\ctypes.CDLL(os.getenv('ROCM_PATH')+'/lib/libhsa-runtime64.so' if os.getenv('ROCM_PATH') else ctypes.util.find_library('hsa-runtime64'))\g" $BASE/hsa.py
python3 -c "import tinygrad.runtime.autogen.hsa"
}
generate_io_uring() {
clang2py -k cdefstum \
/usr/include/liburing.h \
/usr/include/linux/io_uring.h \
-o $BASE/io_uring.py
sed -r '/^#define __NR_io_uring/ s/^#define __(NR_io_uring[^ ]+) (.*)$/\1 = \2/; t; d' /usr/include/asm-generic/unistd.h >> $BASE/io_uring.py # io_uring syscalls numbers
fixup $BASE/io_uring.py
}
generate_ib() {
clang2py -k cdefstum \
/usr/include/infiniband/verbs.h \
/usr/include/infiniband/verbs_api.h \
/usr/include/infiniband/ib_user_ioctl_verbs.h \
/usr/include/rdma/ib_user_verbs.h \
-o $BASE/ib.py
sed -i "s\import ctypes\import ctypes, ctypes.util\g" "$BASE/ib.py"
sed -i "s\FIXME_STUB\libibverbs\g" "$BASE/ib.py"
sed -i "s\FunctionFactoryStub()\ctypes.CDLL(ctypes.util.find_library('ibverbs'), use_errno=True)\g" "$BASE/ib.py"
fixup $BASE/ib.py
}
generate_llvm() {
INC="$(llvm-config-14 --includedir)"
clang2py -k cdefstum \
$(find "$INC/llvm-c/" -type f -name '*.h' | sort) \
"$INC/llvm/Config/Targets.def" \
"$INC/llvm/Config/AsmPrinters.def" \
"$INC/llvm/Config/AsmParsers.def" \
"$INC/llvm/Config/Disassemblers.def" \
--clang-args="$(llvm-config-14 --cflags)" \
-o "$BASE/llvm.py"
sed -i "s\import ctypes\import ctypes, tinygrad.runtime.support.llvm as llvm_support\g" "$BASE/llvm.py"
sed -i "s\FIXME_STUB\llvm\g" "$BASE/llvm.py"
sed -i "s\FunctionFactoryStub()\ctypes.CDLL(llvm_support.LLVM_PATH)\g" "$BASE/llvm.py"
fixup "$BASE/llvm.py"
}
generate_kgsl() {
clang2py extra/qcom_gpu_driver/msm_kgsl.h -o $BASE/kgsl.py -k cdefstum
fixup $BASE/kgsl.py
sed -i "s\import ctypes\import ctypes, os\g" $BASE/kgsl.py
sed -nE 's/#define ([A-Za-z0-9_]+)_SHIFT\s*[^\S\r\n]*[0-9]*$/def \1(val): return (val << \1_SHIFT) \& \1_MASK/p' extra/qcom_gpu_driver/msm_kgsl.h >> $BASE/kgsl.py
sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\__fd.ioctl((__idir<<30)\g" $BASE/kgsl.py
python3 -c "import tinygrad.runtime.autogen.kgsl"
}
generate_adreno() {
clang2py extra/qcom_gpu_driver/a6xx.xml.h -o $BASE/adreno.py -k cestum
sed -nE 's/#define ([A-Za-z0-9_]+)__SHIFT\s*[^\S\r\n]*[0-9]*$/def \1(val): return (val << \1__SHIFT) \& \1__MASK/p' extra/qcom_gpu_driver/a6xx.xml.h >> $BASE/adreno.py
fixup $BASE/adreno.py
sed -i "s\import ctypes\import ctypes, os\g" $BASE/adreno.py
python3 -c "import tinygrad.runtime.autogen.adreno"
}
generate_qcom() {
clang2py -k cdefstum \
extra/dsp/include/ion.h \
extra/dsp/include/msm_ion.h \
extra/dsp/include/adsprpc_shared.h \
extra/dsp/include/remote_default.h \
extra/dsp/include/apps_std.h \
-o $BASE/qcom_dsp.py
fixup $BASE/qcom_dsp.py
python3 -c "import tinygrad.runtime.autogen.qcom_dsp"
}
generate_pci() {
clang2py -k cdefstum \
/usr/include/linux/pci_regs.h \
-o $BASE/pci.py
fixup $BASE/pci.py
}
generate_vfio() {
clang2py -k cdefstum \
/usr/include/linux/vfio.h \
-o $BASE/vfio.py
fixup $BASE/vfio.py
sed -i "s\import ctypes\import ctypes, os\g" $BASE/vfio.py
sed -i "s\import fcntl, functools\import functools" $BASE/vfio.py
sed -i "s\import ctypes,os\a from tinygrad.runtime.support import FileIOInterface\g" $BASE/vfio.py
sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\return __fd.ioctl((__idir<<30)\g" $BASE/vfio.py
}
generate_am() {
AMKERN_COMMIT_HASH=ceb12c04e2b5b53ec0779362831f5ee40c4921e4
AMKERN_SRC=/tmp/ROCK-Kernel-Driver-$AMKERN_COMMIT_HASH
if [ ! -d "$AMKERN_SRC" ]; then
git clone https://github.com/ROCm/ROCK-Kernel-Driver $AMKERN_SRC --depth 1
fi
AMKERN_AMD=$AMKERN_SRC/drivers/gpu/drm/amd/
AMKERN_INC=$AMKERN_AMD/include/
clang2py -k cdefstum \
extra/amdpci/headers/v11_structs.h \
extra/amdpci/headers/v12_structs.h \
extra/amdpci/headers/amdgpu_vm.h \
extra/amdpci/headers/discovery.h \
extra/amdpci/headers/amdgpu_ucode.h \
extra/amdpci/headers/psp_gfx_if.h \
extra/amdpci/headers/amdgpu_psp.h \
extra/amdpci/headers/amdgpu_irq.h \
extra/amdpci/headers/amdgpu_doorbell.h \
$AMKERN_INC/soc15_ih_clientid.h \
--clang-args="-include stdint.h" \
-o $BASE/am/am.py
fixup $BASE/am/am.py
sed -i "s\(int64_t)\ \g" $BASE/am/am.py
sed -i "s\AMDGPU_PTE_MTYPE_VG10(2)\AMDGPU_PTE_MTYPE_VG10(0, 2)\g" $BASE/am/am.py # incorrect parsing (TODO: remove when clang2py is gone).
clang2py -k cdefstum \
$AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
$AMKERN_AMD/amdgpu/soc15d.h \
-o $BASE/am/pm4_soc15.py
fixup $BASE/am/pm4_soc15.py
clang2py -k cdefstum \
$AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
$AMKERN_AMD/amdgpu/nvd.h \
-o $BASE/am/pm4_nv.py
fixup $BASE/am/pm4_nv.py
clang2py -k cdefstum \
extra/hip_gpu_driver/sdma_registers.h \
$AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \
--clang-args="-I/opt/rocm/include -x c++" \
-o $BASE/am/sdma_4_0_0.py
fixup $BASE/am/sdma_4_0_0.py
clang2py -k cdefstum \
extra/hip_gpu_driver/sdma_registers.h \
$AMKERN_AMD/amdgpu/navi10_sdma_pkt_open.h \
--clang-args="-I/opt/rocm/include -x c++" \
-o $BASE/am/sdma_5_0_0.py
fixup $BASE/am/sdma_5_0_0.py
clang2py -k cdefstum \
extra/hip_gpu_driver/sdma_registers.h \
$AMKERN_AMD/amdgpu/sdma_v6_0_0_pkt_open.h \
--clang-args="-I/opt/rocm/include -x c++" \
-o $BASE/am/sdma_6_0_0.py
fixup $BASE/am/sdma_6_0_0.py
clang2py -k cdefstum \
$AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v13_0_0_ppsmc.h \
$AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu13_driver_if_v13_0_0.h \
extra/amdpci/headers/amdgpu_smu.h \
-o $BASE/am/smu_v13_0_0.py
fixup $BASE/am/smu_v13_0_0.py
clang2py -k cdefstum \
$AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_0_pmfw.h \
$AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_2_ppsmc.h \
$AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu14_driver_if_v14_0.h \
extra/amdpci/headers/amdgpu_smu.h \
--clang-args="-include stdint.h" \
-o $BASE/am/smu_v14_0_2.py
fixup $BASE/am/smu_v14_0_2.py
}
generate_sqtt() {
clang2py -k cdefstum \
extra/sqtt/sqtt.h \
-o $BASE/sqtt.py
fixup $BASE/sqtt.py
sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
python3 -c "import tinygrad.runtime.autogen.sqtt"
ROCPROF_COMMIT_HASH=dd0485100971522cc4cd8ae136bdda431061a04d
ROCPROF_SRC=/tmp/rocprof-trace-decoder-$ROCPROF_COMMIT_HASH
if [ ! -d "$ROCPROF_SRC" ]; then
git clone https://github.com/ROCm/rocprof-trace-decoder $ROCPROF_SRC
pushd .
cd $ROCPROF_SRC
git reset --hard $ROCPROF_COMMIT_HASH
popd
fi
clang2py -k cdefstum \
$ROCPROF_SRC/include/rocprof_trace_decoder.h \
$ROCPROF_SRC/include/trace_decoder_instrument.h \
$ROCPROF_SRC/include/trace_decoder_types.h \
-o $BASE/rocprof.py
fixup $BASE/rocprof.py
sed -i '1s/^/# pylint: skip-file\n/' $BASE/rocprof.py
sed -i "s/import ctypes/import ctypes, ctypes.util/g" $BASE/rocprof.py
patch_dlopen $BASE/rocprof.py rocprof-trace-decoder "'/usr/local/lib/librocprof-trace-decoder.so'" "'/usr/local/lib/librocprof-trace-decoder.dylib'"
sed -i "s/def _try_dlopen_rocprof-trace-decoder():/def _try_dlopen_rocprof_trace_decoder():/g" $BASE/rocprof.py
sed -i "s|FunctionFactoryStub()|_try_dlopen_rocprof_trace_decoder()|g" $BASE/rocprof.py
}
generate_webgpu() {
clang2py extra/webgpu/webgpu.h -o $BASE/webgpu.py
fixup $BASE/webgpu.py
sed -i "s/FIXME_STUB/webgpu/g" "$BASE/webgpu.py"
sed -i "s/FunctionFactoryStub()/ctypes.CDLL(webgpu_support.WEBGPU_PATH)/g" "$BASE/webgpu.py"
sed -i "s/import ctypes/import ctypes, tinygrad.runtime.support.webgpu as webgpu_support/g" "$BASE/webgpu.py"
python3 -c "import tinygrad.runtime.autogen.webgpu"
}
generate_libusb() {
clang2py -k cdefstum \
/usr/include/libusb-1.0/libusb.h \
-o $BASE/libusb.py
fixup $BASE/libusb.py
sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libusb.py
sed -i "s/FIXME_STUB/libusb/g" "$BASE/libusb.py"
sed -i "s/libusb_le16_to_cpu = libusb_cpu_to_le16//g" "$BASE/libusb.py"
sed -i "s/FunctionFactoryStub()/None if (lib_path:=os.getenv('LIBUSB_PATH', ctypes.util.find_library('usb-1.0'))) is None else ctypes.CDLL(lib_path)/g" "$BASE/libusb.py"
python3 -c "import tinygrad.runtime.autogen.libusb"
}
generate_mesa() {
MESA_TAG="mesa-25.2.4"
MESA_SRC=/tmp/mesa-$MESA_TAG
TINYMESA_TAG=tinymesa-32dc66c
TINYMESA_DIR=/tmp/tinymesa-$MESA_TAG-$TINYMESA_TAG/
TINYMESA_SO=$TINYMESA_DIR/libtinymesa_cpu.so
if [ ! -d "$MESA_SRC" ]; then
git clone --depth 1 --branch $MESA_TAG https://gitlab.freedesktop.org/mesa/mesa.git $MESA_SRC
pushd .
cd $MESA_SRC
git reset --hard $MESA_COMMIT_HASH
# clang 14 doesn't support packed enums
sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/headers/nv_device_info.h
sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/compiler/nak.h
sed -i "s/nir_instr_type \(\w\+\);/uint8_t \1;/" $MESA_SRC/src/compiler/nir/nir.h
mkdir -p gen/util/format
python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
python3 src/compiler/builtin_types_h.py gen/builtin_types.h
popd
fi
if [ ! -d "$TINYMESA_DIR" ]; then
mkdir $TINYMESA_DIR
curl -L https://github.com/sirhcm/tinymesa/releases/download/$TINYMESA_TAG/libtinymesa_cpu-$MESA_TAG-linux-amd64.so -o $TINYMESA_SO
fi
clang2py -k cdefstu \
$MESA_SRC/src/compiler/nir/nir.h \
$MESA_SRC/src/compiler/nir/nir_builder.h \
$MESA_SRC/src/compiler/nir/nir_shader_compiler_options.h \
$MESA_SRC/src/compiler/nir/nir_serialize.h \
$MESA_SRC/gen/nir_intrinsics.h \
$MESA_SRC/src/nouveau/headers/nv_device_info.h \
$MESA_SRC/src/nouveau/compiler/nak.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_passmgr.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_misc.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_type.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_init.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_nir.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_struct.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_jit_types.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_flow.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_const.h \
$MESA_SRC/src/compiler/glsl_types.h \
$MESA_SRC/src/util/blob.h \
$MESA_SRC/src/util/ralloc.h \
--clang-args="-DHAVE_ENDIAN_H -DHAVE_STRUCT_TIMESPEC -DHAVE_PTHREAD -I$MESA_SRC/src -I$MESA_SRC/include -I$MESA_SRC/gen -I$MESA_SRC/src/compiler/nir -I$MESA_SRC/src/gallium/auxiliary -I$MESA_SRC/src/gallium/include -I$(llvm-config-20 --includedir)" \
-l $TINYMESA_SO \
-o $BASE/mesa.py
LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)
fixup $BASE/mesa.py
patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "'/opt/homebrew/lib/libtinymesa_cpu.dylib'" "'/opt/homebrew/lib/libtinymesa.dylib'"
echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
sed -i "/in_dll/s/.*/try: &\nexcept (AttributeError, ValueError): pass/" $BASE/mesa.py
sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
echo "def __getattr__(nm): raise AttributeError('LLVMpipe requires tinymesa_cpu' if 'tinymesa_cpu' not in dll._name else f'attribute {nm} not found') if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
sed -i "s/ctypes.glsl_base_type/glsl_base_type/" $BASE/mesa.py
# bitfield bug in clang2py
sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
sed -i '/struct_nir_builder._pack_ = 1 # source:False/d' "$BASE/mesa.py"
python3 -c "import tinygrad.runtime.autogen.mesa"
}
if [ "$1" == "opencl" ]; then generate_opencl
elif [ "$1" == "hip" ]; then generate_hip
elif [ "$1" == "comgr" ]; then generate_comgr
elif [ "$1" == "cuda" ]; then generate_cuda
elif [ "$1" == "nvrtc" ]; then generate_nvrtc
elif [ "$1" == "hsa" ]; then generate_hsa
elif [ "$1" == "kfd" ]; then generate_kfd
elif [ "$1" == "nv" ]; then generate_nv
elif [ "$1" == "amd" ]; then generate_amd
elif [ "$1" == "am" ]; then generate_am
elif [ "$1" == "sqtt" ]; then generate_sqtt
elif [ "$1" == "qcom" ]; then generate_qcom
elif [ "$1" == "io_uring" ]; then generate_io_uring
elif [ "$1" == "ib" ]; then generate_ib
elif [ "$1" == "llvm" ]; then generate_llvm
elif [ "$1" == "kgsl" ]; then generate_kgsl
elif [ "$1" == "adreno" ]; then generate_adreno
elif [ "$1" == "pci" ]; then generate_pci
elif [ "$1" == "vfio" ]; then generate_vfio
elif [ "$1" == "webgpu" ]; then generate_webgpu
elif [ "$1" == "libusb" ]; then generate_libusb
elif [ "$1" == "mesa" ]; then generate_mesa
elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_am; generate_webgpu; generate_mesa
else echo "usage: $0 <type>"
fi
examples/openpilot/compile3.py
+1
View File
@@ -1,5 +1,6 @@
import os, sys, pickle, time, re
import numpy as np
if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
from tinygrad.helpers import DEBUG, getenv
extra/bandwidth_test.py
+34
View File
@@ -0,0 +1,34 @@
#!/usr/bin/env python3
from tinygrad import Tensor, Device, GlobalCounters, Context, dtypes
from tinygrad.helpers import getenv, colored
SZ = 8_000_000_000
GPUS = getenv("GPUS", 4) # TODO: expose a way in tinygrad to access this
if __name__ == "__main__":
# create tensors
tens = [Tensor.ones(SZ, dtype=dtypes.uint8, device=f"{Device.DEFAULT}:{i}").contiguous() for i in range(GPUS)]
Tensor.realize(*tens)
bw = [[0.0]*GPUS for _ in range(GPUS)]
for i in range(GPUS):
for j in range(GPUS):
GlobalCounters.reset()
with Context(DEBUG=2):
if i == j:
# this copy would be optimized out, just add 1
(tens[i]+1).realize()
else:
tens[i].to(f"{Device.DEFAULT}:{j}").realize()
t = max(GlobalCounters.time_sum_s, 1e-9)
bw[i][j] = SZ / t / 1e9 # GB/s
def fmt(x):
c = "green" if x > 50 else "yellow" if x > 20 else "red"
return colored(f"{x:6.1f}", c)
# header
print(" " * 8 + " ".join(f"{'d'+str(j):>6}" for j in range(GPUS)))
# rows
for i in range(GPUS):
print(f"{'s'+str(i):>6} -> " + " ".join(fmt(x) for x in bw[i]))
extra/gemm/mi350x_uop_matmul_2.py
+141
View File
@@ -0,0 +1,141 @@
import os
import numpy as np
np.set_printoptions(linewidth=1000000)
os.environ["AMD_LLVM"] = "0"
from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
from tinygrad.helpers import DEBUG, getenv
from tinygrad.dtype import AddrSpace
from tinygrad.uop.ops import sint, AxisType, KernelInfo, Ops
WARP_SIZE = 64
# Reg tile sizes (tensor cores)
TC_M = 16
TC_N = 16
TC_K = 32
N,M,K = 4096,4096,4096
# Threadblock tile sizes (block-level tile of C that a block computes)
BLOCK_M = 64
BLOCK_N = 64
BLOCK_K = 64
WARPGROUP_SIZE = 1
BLOCK_M = BLOCK_M * WARPGROUP_SIZE
TID_SIZE = WARPGROUP_SIZE*WARP_SIZE
def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
assert dest.shape == src.shape
rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
copy = dest[*rngs].store(src[*rngs]).end(*rngs)
return dest.after(copy) if set else copy
def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...], warpgroup, warp) -> UOp:
K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
# load from locals into registers
Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)
M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
Asl = Asl.reshape(BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M)
load_rng = UOp.range(8, rng+11, axis_type=AxisType.UPCAST)
A_in = Asl[K_inner_loop, (warp//16)*8+load_rng, M_load_loop, warpgroup, warp%16].contract(load_rng)
Ar = Ar[M_load_loop].set(A_in, end=M_load_loop)
N_load_loop = UOp.range(BLOCK_N//TC_N, rng+20)
Bsl = Bsl.reshape(BLOCK_K//TC_K, TC_K, BLOCK_N//TC_N, TC_N)
load_rng = UOp.range(8, rng+21, axis_type=AxisType.UPCAST)
B_in = Bsl[K_inner_loop, (warp//16)*8+load_rng, N_load_loop, warp%16].contract(load_rng)
Br = Br[N_load_loop].set(B_in, end=N_load_loop)
M_inner_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+30)
N_inner_loop = UOp.range(BLOCK_N//TC_N, rng+31)
# load values
acc_after = acc.after(*afters, M_inner_loop, N_inner_loop, K_inner_loop)
acc_load = acc_after[N_inner_loop, M_inner_loop]
# do WMMA
wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
# store back the acc
acc_store = acc[N_inner_loop, M_inner_loop].store(out)
return acc_store.end(M_inner_loop, N_inner_loop, K_inner_loop)
def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
gx, gy = UOp.special(M//BLOCK_M, "gidx0"), UOp.special(N//BLOCK_N, "gidx1")
K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)
# split out the globals into blocks
C = C.src[0].cast(dtypes.float.vec(4).ptr(C.ptrdtype.size)).reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
A = A.reshape((M//BLOCK_M, BLOCK_M, K//BLOCK_K, BLOCK_K))[gx, :, K_outer_loop, :]
B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))[K_outer_loop, :, gy, :]
# ---------------------------
# GLOBAL -> LOCAL (As, Bs)
# ---------------------------
tid = UOp.special(TID_SIZE, "lidx0")
warpgroup, warp = tid//WARP_SIZE, tid%WARP_SIZE
A_view = A.reshape(-1, TID_SIZE, 8)
B_view = B.reshape(-1, TID_SIZE, 8)
# A: read BM x BK tiles (permute on store into locals)
As = UOp.placeholder((BLOCK_K, BLOCK_M), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL).shrink_to(BLOCK_K, BLOCK_M)
As_view = As.reshape(-1, TID_SIZE, 8)
Bs = UOp.placeholder((BLOCK_K, BLOCK_N+4), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL).shrink_to(BLOCK_K, BLOCK_N)
Bs_view = Bs.reshape(-1, TID_SIZE, 8)
outer_copy = UOp.range(A_view.shape[0], 100, AxisType.UPCAST)
inner_copy = UOp.range(A_view.shape[2], 101, AxisType.UPCAST)
As_store = As_view[outer_copy, tid, inner_copy].store(A_view[outer_copy, tid, inner_copy])
Bs_store = Bs_view[outer_copy, tid, inner_copy].store(B_view[outer_copy, tid, inner_copy])
if getenv("NOLOAD"):
As_store = As[0,0].store(0)
Bs_store = Bs[0,0].store(0)
# TODO: can we automate barrier?
barrier = UOp.barrier(UOp.group(As_store, Bs_store).end(outer_copy, inner_copy))
if getenv("COMPUTE"):
As, Bs = As.after(barrier), Bs.after(barrier)
acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)
sink = compute_on_locals(acc, As, Bs, 200, afters=(barrier,), warpgroup=warpgroup, warp=warp)
sink = sink.end(K_outer_loop)
C_view = C[gx, :, gy, :].reshape(BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M, BLOCK_N//TC_N, TC_N)[:, warpgroup, warp%16, :, (warp//16)*4]
sink = copy(C_view, acc.after(sink), rng=300)
else:
sink = C.after(barrier.end(K_outer_loop))[0,0,0,0].store(As[0,0]+Bs[0,0])
return sink.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
if __name__ == "__main__":
a = Tensor.randn(M, K, dtype=dtypes.half)
b = Tensor.randn(K, N, dtype=dtypes.half)
c = Tensor.empty(M, N, dtype=dtypes.float)
with Context(DEBUG=0): Tensor.realize(a,b)
GlobalCounters.reset()
with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
tst.realize()
print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
with Context(DEBUG=0):
ref = a.dot(b, dtype=dtypes.float)
ref.realize()
#print(ref.numpy())
#print(tst.numpy())
assert Tensor.isclose(ref, tst, atol=1e-2).all().item(), "matrix not close"
extra/gemm/torch_gemm.py
+1 -1
View File
@@ -12,7 +12,7 @@ MPS = getenv("MPS", 0)
if getenv("FP16_ACC"): torch.backends.cuda.matmul.allow_fp16_accumulation = True
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
for N in [256, 512, 1024, 2048, 4096]:
for N in [256, 512, 1024, 2048, 4096] + ([6144, 8192] if getenv("BIG") else []):
FLOPS = N*N*N*2
b = torch.rand((N,N), dtype=dtype)
extra/gpuburn.py
+16
View File
@@ -0,0 +1,16 @@
from tinygrad import Tensor, Device, TinyJit, dtypes
from tinygrad.helpers import getenv
GPUS = getenv("GPUS", 4) # TODO: expose a way in tinygrad to access this
N = 6144
@TinyJit
def many_matmul(A, B):
out = A
for _ in range(8): out = out@B
return out
if __name__ == "__main__":
A = Tensor.ones(GPUS, N, N, dtype=dtypes.half).shard(devices=tuple([f"{Device.DEFAULT}:{i}" for i in range(GPUS)]), axis=0).contiguous()
B = Tensor.ones(GPUS, N, N, dtype=dtypes.half).shard(devices=tuple([f"{Device.DEFAULT}:{i}" for i in range(GPUS)]), axis=0).contiguous()
while 1: many_matmul(A, B)
extra/hcqfuzz/tools.py
+4
View File
@@ -51,11 +51,15 @@ def create_report(dev, test, result, stdout, stderr):
dmesg_output = subprocess.check_output(["sudo", "dmesg", "--ctime", "--color=never"], text=True)
with open(dmesg_path, "w") as f: f.write(dmesg_output)
env_vars = " ".join(f"{k}={v}" for k, v in test.env.items())
reproduce_cmd = f"{env_vars} {test.cmd}"
summary_path = os.path.join(report_path, "summary.txt")
with open(summary_path, "w") as f:
f.write(f"Test: {test.name()}\n")
f.write(f"Dev params: {vars(dev)}\n")
f.write(f"Test params: {vars(test)}\n")
f.write(f"Reproduce cmd: {reproduce_cmd}\n")
f.write(f"Exit Code: {result}\n")
print(f"Crash report saved to {report_path}")
extra/mesa/lvp_nir_options.sh
+2 -1
View File
@@ -19,5 +19,6 @@ trap 'rm -f "$TMP"' EXIT
EOF
sed -n '/struct nir_shader_compiler_options/,/^}/{p;/^}/q}' $1/src/gallium/drivers/llvmpipe/lp_screen.c
echo "int main(void) { write(1, &gallivm_nir_options, sizeof(gallivm_nir_options)); }"
) | cc -x c -o $TMP - -I$1/src/compiler/nir -I$1/src -I$1/include && $TMP | gzip | base64 -w0
) | cc -x c -o $TMP - -I$1/src/compiler/nir -I$1/src -I$1/include || exit 1
printf 'lvp_nir_options = gzip.decompress(base64.b64decode("%s"))' $("$TMP" | gzip | base64 -w0)
extra/mmapeak/mmapeak.py
+1 -1
View File
@@ -32,7 +32,7 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, extra="")
src = src.replace("DIRECTIVE", DIRECTIVE)
lib = COMPILER.compile(src)
fxn = AMDProgram(DEV, "matmul", lib)
elapsed = fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True)
elapsed = min([fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) for _ in range(2)])
FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
print(f"{instruction:<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")
extra/mmapeak/template.s
+8 -8
View File
@@ -3,14 +3,14 @@
.p2align 8
.type matmul,@function
matmul:
s_mov_b32 s1, INTERNAL_LOOP
s_mov_b32 s2, 0
inner_loop:
INSTRUCTION
s_sub_u32 s1, s1, 1
s_cmp_lg_i32 s1, s2
s_cbranch_scc1 inner_loop
s_endpgm
s_mov_b32 s1, INTERNAL_LOOP
s_mov_b32 s2, 0
inner_loop:
INSTRUCTION
s_sub_u32 s1, s1, 1
s_cmp_lg_i32 s1, s2
s_cbranch_scc1 inner_loop
s_endpgm
.rodata
.p2align 6
extra/nvJitLink.h
+53
View File
@@ -0,0 +1,53 @@
/*
* NVIDIA_COPYRIGHT_BEGIN
*
* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
*
* NVIDIA CORPORATION and its licensors retain all intellectual property
* and proprietary rights in and to this software, related documentation
* and any modifications thereto. Any use, reproduction, disclosure or
* distribution of this software and related documentation without an express
* license agreement from NVIDIA CORPORATION is strictly prohibited.
*
* NVIDIA_COPYRIGHT_END
*/
#include <stdint.h>
#include <stdlib.h>
typedef enum {
NVJITLINK_SUCCESS = 0,
NVJITLINK_ERROR_UNRECOGNIZED_OPTION,
NVJITLINK_ERROR_MISSING_ARCH,
NVJITLINK_ERROR_INVALID_INPUT,
NVJITLINK_ERROR_PTX_COMPILE,
NVJITLINK_ERROR_NVVM_COMPILE,
NVJITLINK_ERROR_INTERNAL
} nvJitLinkResult;
typedef enum {
NVJITLINK_INPUT_NONE = 0,
NVJITLINK_INPUT_CUBIN = 1,
NVJITLINK_INPUT_PTX,
NVJITLINK_INPUT_LTOIR,
NVJITLINK_INPUT_FATBIN,
NVJITLINK_INPUT_OBJECT,
NVJITLINK_INPUT_LIBRARY
} nvJitLinkInputType;
typedef struct nvJitLink* nvJitLinkHandle;
nvJitLinkResult nvJitLinkCreate(nvJitLinkHandle *handle, uint32_t numOptions, const char **options);
nvJitLinkResult nvJitLinkDestroy(nvJitLinkHandle *handle);
nvJitLinkResult nvJitLinkAddData(nvJitLinkHandle handle, nvJitLinkInputType inputType, const void *data, size_t size, const char *name);
nvJitLinkResult nvJitLinkAddFile(nvJitLinkHandle handle, nvJitLinkInputType inputType, const char *fileName);
nvJitLinkResult nvJitLinkComplete(nvJitLinkHandle handle);
nvJitLinkResult nvJitLinkGetLinkedCubinSize(nvJitLinkHandle handle, size_t *size);
nvJitLinkResult nvJitLinkGetLinkedCubin(nvJitLinkHandle handle, void *cubin);
nvJitLinkResult nvJitLinkGetLinkedPtxSize(nvJitLinkHandle handle, size_t *size);
nvJitLinkResult nvJitLinkGetLinkedPtx(nvJitLinkHandle handle, char *ptx);
nvJitLinkResult nvJitLinkGetErrorLogSize(nvJitLinkHandle handle, size_t *size);
nvJitLinkResult nvJitLinkGetErrorLog(nvJitLinkHandle handle, char *log);
nvJitLinkResult nvJitLinkGetInfoLogSize(nvJitLinkHandle handle, size_t *size);
nvJitLinkResult nvJitLinkGetInfoLog(nvJitLinkHandle handle, char *log);
nvJitLinkResult nvJitLinkVersion(unsigned int *major, unsigned int *minor);
extra/nv_gpu_driver/clcec0qmd.h
+2
View File
@@ -65,6 +65,8 @@
#define NVCEC0_QMDV05_00_GRID_HEIGHT_RESUME MW(271:256)
#define NVCEC0_QMDV05_00_GRID_DEPTH_RESUME MW(287:272)
#define NVCEC0_QMDV05_00_RELEASE_ENABLE(i) MW((288+(i)*16):(288+(i)*16))
#define NVCEC0_QMDV05_00_RELEASE0_ENABLE NVCEC0_QMDV05_00_RELEASE_ENABLE(0)
#define NVCEC0_QMDV05_00_RELEASE1_ENABLE NVCEC0_QMDV05_00_RELEASE_ENABLE(1)
#define NVCEC0_QMDV05_00_RELEASE_ENABLE_FALSE 0x00000000
#define NVCEC0_QMDV05_00_RELEASE_ENABLE_TRUE 0x00000001
#define NVCEC0_QMDV05_00_RELEASE_STRUCTURE_SIZE(i) MW((290+(i)*16):(289+(i)*16))
extra/nv_gpu_driver/nv_ioctl.py
+3 -3
View File
@@ -58,13 +58,13 @@ def install_hook(c_function, python_function):
return orig_func
# *** ioctl lib end ***
import tinygrad.runtime.autogen.nv_gpu as nv_gpu
from tinygrad.runtime.autogen import nv_570 as nv_gpu
nvescs = {getattr(nv_gpu, x):x for x in dir(nv_gpu) if x.startswith("NV_ESC")}
nvcmds = {getattr(nv_gpu, x):(x, getattr(nv_gpu, "struct_"+x+"_PARAMS", getattr(nv_gpu, "struct_"+x.replace("_CMD_", "_")+"_PARAMS", None))) for x in dir(nv_gpu) if \
x.startswith("NV") and x[6:].startswith("_CTRL_") and isinstance(getattr(nv_gpu, x), int)}
def get_classes():
hdrpy = (pathlib.Path(__file__).parent.parent.parent / "tinygrad/runtime/autogen/nv_gpu.py").read_text()
hdrpy = (pathlib.Path(__file__).parent.parent.parent / "tinygrad/runtime/autogen/nv_570.py").read_text()
clss = re.search(r'NV01_ROOT.*?NV_SEMAPHORE_SURFACE = \(0x000000da\) # macro', hdrpy, re.DOTALL).group()
pattern = r'([0-9a-zA-Z_]*) = +\((0x[0-9a-fA-F]+)\)'
matches = re.findall(pattern, clss, re.MULTILINE)
@@ -272,4 +272,4 @@ def compare_launch_state(states, good_states):
return True, "PASS"
# IOCTL=1 CUDA=1 CUDA_PTX=1 python3 test/test_ops.py TestOps.test_tiny_add
# IOCTL=1 CUDA=1 CUDA_PTX=1 python3 test/test_ops.py TestOps.test_tiny_add
examples/self_tokenize.py → extra/self_tokenize.py
+9 -11
View File
@@ -2,7 +2,8 @@ import os, pathlib, argparse
from examples.llama3 import Tokenizer
from tabulate import tabulate
from tinygrad import fetch
from tinygrad.helpers import flatten
from tinygrad.helpers import flatten, getenv
from sz import NONCORE_DIRS
# llama 3 tokenizer
tokenizer = Tokenizer(fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model").as_posix())
@@ -10,19 +11,15 @@ tokenizer = Tokenizer(fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/
def read_code(base_path):
ret = []
for path, _, files in os.walk(os.path.join(base_path, "tinygrad")):
if not getenv("CORE") and any(path.split("./")[1].startswith(x) for x in NONCORE_DIRS): continue
for name in files:
if not name.endswith(".py"): continue
if 'tinygrad/runtime/autogen' in path.replace('\\', '/'): continue
fullpath = os.path.join(path, name)
code = pathlib.Path(fullpath).read_text()
ret.append(("### " + fullpath.split("tinygrad/", 1)[1], code))
ret.append((fullpath.split("tinygrad/", 1)[1], code))
return ret
def write_code_to_file(filename, code_list):
"""Writes the combined code to a specified file."""
with open(filename, 'w') as f:
f.write('\n'.join(flatten(code_list)))
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Analyze and optionally save tinygrad code.")
parser.add_argument("--output", help="Output file to write the combined code to.")
@@ -32,10 +29,11 @@ if __name__ == "__main__":
table = []
for name,code in ret:
table.append([name, len(tokenizer.encode(name+"\x00"+code))])
table.append([name, len(tokenizer.encode(code))])
print(tabulate([["name", "llm tokens"]]+sorted(table, key=lambda x: -x[1]), headers="firstrow"))
code_str = '\x00'.join(flatten(ret))
banner = "#"*40
code_str = ''.join([f"{banner}\n# {name}\n{banner}\n\n{code}\n" for name,code in ret])
print(f"code has {len(code_str)} chars")
newline_count = code_str.count('\n')
print(f"code has {newline_count} newlines")
@@ -44,5 +42,5 @@ if __name__ == "__main__":
print(f"code has {len(encoded)} tokens")
if args.output:
write_code_to_file(args.output, ret)
print(f"Combined code written to {args.output}")
with open(args.output, 'w') as f: f.write(code_str)
print(f"Combined code written to {args.output}")
extra/sqtt/install_sqtt_decoder.py
+4 -2
View File
@@ -12,7 +12,9 @@ if __name__ == "__main__":
lib = fp.parent/"rocprof-trace-decoder-macos-arm64-0.1.4-Darwin"/"lib"/"librocprof-trace-decoder.dylib"
os.chmod(fp, 0o755)
os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir")
shutil.copy2(lib, DEST)
else:
lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/43bf0fef74a83c3c25badfc5a09c0bd39ed8c6f9/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so")
shutil.copy2(lib, DEST)
lib = DEST/"librocprof-trace-decoder.so"
os.system("sudo curl -L https://github.com/ROCm/rocprof-trace-decoder/raw/43bf0fef74a83c3c25badfc5a09c0bd39ed8c6f9/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so -o"+str(lib))
os.system("sudo ldconfig")
print(f"Installed {lib.name} to", DEST)
extra/sqtt/rgptool.py
+7 -11
View File
@@ -185,9 +185,7 @@ class RGP:
magic_number=sqtt.SQTT_FILE_MAGIC_NUMBER,
version_major=sqtt.SQTT_FILE_VERSION_MAJOR,
version_minor=sqtt.SQTT_FILE_VERSION_MINOR,
flags=sqtt.struct_sqtt_file_header_flags(
_0=sqtt.union_sqtt_file_header_flags_0(value=1),
),
flags=sqtt.struct_sqtt_file_header_flags(value=1,),
chunk_offset=ctypes.sizeof(sqtt.struct_sqtt_file_header),
)
chunks = [
@@ -265,7 +263,7 @@ class RGP:
profiling_mode=sqtt.SQTT_PROFILING_MODE_PRESENT,
instruction_trace_mode=sqtt.SQTT_INSTRUCTION_TRACE_FULL_FRAME if sqtt_itrace_enabled else sqtt.SQTT_INSTRUCTION_TRACE_DISABLED,
instruction_trace_data=sqtt.union_sqtt_instruction_trace_data(
shader_engine_filter=sqtt.struct_sqtt_instruction_trace_data_shader_engine_filter(mask=sqtt_itrace_se_mask),
shader_engine_filter=sqtt.union_sqtt_instruction_trace_data_shader_engine_filter(mask=sqtt_itrace_se_mask),
),
)),
*flatten([(
@@ -276,13 +274,11 @@ class RGP:
),
shader_engine_index=sqtt_event.se,
sqtt_version={11: sqtt.SQTT_VERSION_3_2, 12: sqtt.SQTT_VERSION_3_3}.get(gfx_ver),
_0=sqtt.union_sqtt_file_chunk_sqtt_desc_0(
v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_0_v1(
instrumentation_spec_version=1,
instrumentation_api_version=0,
compute_unit_index=0,
)
),
v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_0_v1(
instrumentation_spec_version=1,
instrumentation_api_version=0,
compute_unit_index=0,
)
)),
RGPChunk(sqtt.struct_sqtt_file_chunk_sqtt_data(
header=sqtt.struct_sqtt_file_chunk_header(
extra/sqtt/roc.py
+13 -35
View File
@@ -28,18 +28,6 @@ def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
cur_off += instr_sz
return addr_table
@dataclasses.dataclass
class InstInfo:
typ:str=""
inst:str=""
hit:int=0
lat:int=0
stall:int=0
def __str__(self): return f"{self.inst:>20} hits:{self.typ:>6} hits:{self.hit:>6} latency:{self.lat:>6} stall:{self.stall:>6}"
def on_ev(self, ev):
self.hit, self.lat, self.stall = self.hit + 1, self.lat + ev.duration, self.stall + ev.stall
@dataclasses.dataclass(frozen=True)
class InstExec:
typ:str
@@ -48,25 +36,18 @@ class InstExec:
dur:int
time:int
@dataclasses.dataclass(frozen=True)
class PrgExec:
name:str
wave:int
cu:int
simd:int
def __str__(self): return f"{self.name},{self.wave},{self.cu},{self.simd}"
@dataclasses.dataclass(frozen=True)
class WaveExec:
wave_id:int
cu:int
simd:int
begin_time:int
end_time:int
insts:list[InstExec]
class _ROCParseCtx:
def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
self.dev_evs, self.sqtt_evs, self.prog_evs = dev_evs, iter(sqtt_evs), prog_evs
self.wave_events:dict[PrgExec, dict[int, InstInfo]] = {}
self.disasms:dict[tuple[str, int], tuple[str, int]] = {}
self.inst_execs:dict[str, list[WaveExec]] = {}
@@ -79,7 +60,8 @@ class _ROCParseCtx:
x = next(self.sqtt_evs, None)
self.active_kern = x.kern if x is not None else None
self.active_se = x.se if x is not None else None
return x
self.active_blob = (ctypes.c_ubyte * len(x.blob)).from_buffer_copy(x.blob) if x is not None else None
return self.active_blob
def on_occupancy_ev(self, ev):
if DEBUG >= 5: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
@@ -87,19 +69,15 @@ class _ROCParseCtx:
def on_wave_ev(self, ev):
if DEBUG >= 5: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
asm:dict[int, InstInfo] = {}
inst_execs:list[InstExec] = []
for j in range(ev.instructions_size):
inst_ev = ev.instructions_array[j]
inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
inst_typ = rocprof.enum_rocprofiler_thread_trace_decoder_inst_category_t.get(inst_ev.category)
inst_disasm = self.disasms[(unwrap(self.active_kern), unwrap(inst_ev.pc.address))][0]
asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=inst_disasm))
asm[inst_ev.pc.address].on_ev(inst_ev)
inst_execs.append(InstExec(inst_typ, inst_disasm, inst_ev.stall, inst_ev.duration, inst_ev.time))
if ev.instructions_size > 0:
self.wave_events[key:=PrgExec(unwrap(self.active_kern), ev.wave_id, ev.cu, ev.simd)] = asm
self.inst_execs.setdefault(key.name, []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, inst_execs))
self.inst_execs.setdefault(unwrap(self.active_kern), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, ev.begin_time, ev.end_time, inst_execs))
def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
dev_events:dict[str, ProfileDeviceEvent] = {}
@@ -114,10 +92,10 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
@rocprof.rocprof_trace_decoder_se_data_callback_t
def copy_cb(buf, buf_size, data_ptr):
if (prof:=ROCParseCtx.next_sqtt()) is None: return 0
buf[0] = ctypes.cast((ctypes.c_ubyte * len(prof.blob)).from_buffer_copy(prof.blob), ctypes.POINTER(ctypes.c_ubyte))
buf_size[0] = len(prof.blob)
return len(prof.blob)
if (prof_info:=ROCParseCtx.next_sqtt()) is None: return 0
buf[0] = ctypes.cast(prof_info, ctypes.POINTER(ctypes.c_ubyte))
buf_size[0] = len(prof_info)
return len(prof_info)
@rocprof.rocprof_trace_decoder_trace_callback_t
def trace_cb(record_type, events_ptr, n, data_ptr):
@@ -127,7 +105,7 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr): ROCParseCtx.on_wave_ev(ev)
case _:
if DEBUG >= 5: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
if DEBUG >= 5: print(rocprof.enum_rocprofiler_thread_trace_decoder_record_type_t.get(record_type), events_ptr, n)
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
@rocprof.rocprof_trace_decoder_isa_callback_t
@@ -147,7 +125,7 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
try:
rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run ./extra/sqtt/install_sqtt_decoder.py to install") from e
except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
return ROCParseCtx
if __name__ == "__main__":
@@ -157,7 +135,7 @@ if __name__ == "__main__":
with args.profile.open("rb") as f: profile = pickle.load(f)
rctx = decode(profile)
print('SQTT:', rctx.wave_events.keys())
print('SQTT:', rctx.inst_execs.keys())
for ev in profile:
if not isinstance(ev, ProfilePMCEvent): continue
extra/sqtt/test_timing.py
+26 -19
View File
@@ -1,7 +1,7 @@
import os
os.environ["PYTHONPATH"] = "."
os.environ["SQTT"] = "1"
os.environ["AMD"] = "1"
if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
os.environ["VIZ"] = "1"
os.environ["AMD_LLVM"] = "0"
@@ -10,13 +10,13 @@ import sys, contextlib
from tinygrad import Tensor
from tinygrad.dtype import dtypes
from tinygrad.renderer import ProgramSpec
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AddrSpace
from tinygrad.engine.realize import CompiledRunner
from tinygrad.device import Device, ProfileDeviceEvent
from extra.sqtt.roc import decode, InstExec, PrgExec
from extra.sqtt.roc import decode, WaveExec
dev = Device["AMD"]
dev = Device[os.environ["DEV"]]
def custom(arg:str, s:UOp|None=None) -> UOp: return UOp(Ops.CUSTOM, src=(s,) if s is not None else (), arg=arg)
@@ -36,12 +36,13 @@ def asm_kernel(instrs:list[str], l:int=1, g:int=1) -> Tensor:
def save_sqtt():
# clear the old traces
dev.profile_events.clear()
sqtt:dict[PrgExec, list[InstExec]] = {}
sqtt:dict[str, list[WaveExec]] = {}
yield sqtt
# decode sqtt
rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
assert len(rctx.inst_execs) > 0, "empty sqtt output"
sqtt.update(rctx.inst_execs)
if os.environ["DEV"] == "AMD":
rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
assert len(rctx.inst_execs) > 0, "empty sqtt output"
sqtt.update(rctx.inst_execs)
class TestTiming(unittest.TestCase):
def test_v_add(self):
@@ -62,18 +63,24 @@ class TestTiming(unittest.TestCase):
assert all(s.stall == 0 for s in wave)
def test_multi_cycle_inst(self):
def custom_vrcp(A, B):
op = custom("float a = 0.0;")
op = custom("float b = (*(data1_1+0));", op)
#op = custom('asm volatile("v_mul_f32_e32 %2 %2 %1" : "+v"(a) : "v"(b));', op)
op = custom('asm volatile("v_rcp_f32_e32 %2 %1" : "+v"(a) : "v"(b));', op)
op = custom('asm volatile("v_add_f32_e64 %1 %1 1.0" : "+v"(a));', op)
op = custom("*(data0_1+0) = a;", op)
return UOp.sink(op, A, B, arg=KernelInfo(name="custom_vrcp"))
out = Tensor([0.]).realize()
inp = Tensor([-2.0]).realize()
with save_sqtt() as sqtt:
asm_kernel([
"v_mov_b32_e32 v4 0x3f800000",
"v_rcp_f32_e32 v5 v4",
"v_mul_f32_e32 v6 v5 v4",
]).realize()
w = list(sqtt.values())[0]
rcp, mul = w[1], w[2]
self.assertGreater(rcp.dur, 1) # 4 cycles on gfx11
self.assertEqual(mul.dur, 1)
# mul depends on v5, how can it run before rcp is done?
self.assertGreaterEqual(mul.time, rcp.time+rcp.dur)
Tensor.custom_kernel(out, inp, fxn=custom_vrcp)[0].realize()
wave = list(sqtt.values())[0][0]
for i in range(len(wave.insts)):
if wave.insts[i].inst.startswith("global_store"):
print(f"store diff {wave.insts[i].time-(wave.insts[i-1].time)}")
self.assertEqual(out.item(), 0.5)
def test_wmma(self):
with save_sqtt() as sqtt:
extra/thunder/tiny/tk/group.py
+104 -117
View File
@@ -7,7 +7,7 @@ from tinygrad.dtype import AddrSpace, PtrDType
from tinygrad.helpers import getenv, prod
from extra.thunder.tiny.tk import WARP_THREADS
from extra.thunder.tiny.tk.tiles import TILE_ROW_DIM, TILE_COL_DIM, RT_BASE_TILE_NEPT, slots
from extra.thunder.tiny.tk.tiles import ALL_TILES, GL, ST, RT, RV
class Group:
def __init__(self, warps:int, ker):
@@ -27,27 +27,26 @@ class Group:
# ops that only work on a single warp
clear_rid = 1000
def clear(self, reg:UOp, value:float=0):
def clear(self, reg:ALL_TILES, value:float=0):
reg = cast(UOp, reg)
assert self.warps == 1
i = UOp.range(reg.size, Group.clear_rid)
Group.clear_rid += 1
rngs_for_shape = tuple(UOp.range(dim, Group.clear_rid + i) for i, dim in enumerate(reg.shape))
Group.clear_rid += len(reg.shape)
reg_store = reg.reshape((reg.size,))[i].store(value).end(i)
reg_store = reg[*rngs_for_shape].store(value).end(*rngs_for_shape)
self.ker.push_store(reg_store, reg)
return reg.after(reg_store).reshape(reg.shape)
def zero(self, reg:UOp): return self.clear(reg, 0)
def neg_inf(self, reg:UOp): return self.clear(reg, -math.inf)
def zero(self, reg:ALL_TILES): return self.clear(reg, 0)
def neg_inf(self, reg:ALL_TILES): return self.clear(reg, -math.inf)
copy_rid = 300
def copy(self, dst:UOp, src:UOp):
def copy(self, dst:ALL_TILES, src:ALL_TILES):
dst, src = cast(UOp, dst), cast(UOp, src)
assert self.warps == 1
assert dst.shape == src.shape
assert cast(PtrDType, dst.dtype).addrspace == AddrSpace.REG
assert cast(PtrDType, src.dtype).addrspace == AddrSpace.REG
rngs_for_shape = tuple(UOp.range(dim, Group.copy_rid + i) for i, dim in enumerate(dst.shape))
Group.copy_rid += len(dst.shape)
@@ -57,57 +56,55 @@ class Group:
self.ker.push_store(dst_store, dst)
return dst.after(dst_store).reshape(dst.shape)
mma_rid = 600
def mma_AB(self, c:UOp, a:UOp, b:UOp, after=True):
def mma_AB(self, c:UOp|RT, a:UOp|RT, b:UOp|RT, after=True):
c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
assert self.warps == 1
mma_i_height = UOp.range(c.shape[-3], Group.mma_rid)
mma_i_width = UOp.range(c.shape[-2], Group.mma_rid+1)
mma_i_inner = UOp.range(a.shape[-2], Group.mma_rid+2, AxisType.REDUCE)
Group.mma_rid += 3
for height in self.ker.range(c.shape[-3], track=False):
for width in self.ker.range(c.shape[-2], track=False):
for inner in self.ker.range(a.shape[-2], AxisType.REDUCE, track=False):
wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
b_in1 = UOp.vectorize(*([b[inner, width, i] for i in range(2)] + [b[inner, width, 4+i] for i in range(2)]))
c_out1 = UOp.vectorize(*[c[height, width, i] for i in range(4)])
b_in2 = UOp.vectorize(*([b[inner, width, 2+i] for i in range(2)] + [b[inner, width, 6+i] for i in range(2)]))
c_out2 = UOp.vectorize(*[c[height, width, 4+i] for i in range(4)])
a_in = UOp.vectorize(*[a[mma_i_height, mma_i_inner, i] for i in range(8)])
b_in1 = UOp.vectorize(*([b[mma_i_inner, mma_i_width, i] for i in range(2)] + [b[mma_i_inner, mma_i_width, 4+i] for i in range(2)]))
c_out1 = UOp.vectorize(*[c[mma_i_height, mma_i_width, i] for i in range(4)])
b_in2 = UOp.vectorize(*([b[mma_i_inner, mma_i_width, 2+i] for i in range(2)] + [b[mma_i_inner, mma_i_width, 6+i] for i in range(2)]))
c_out2 = UOp.vectorize(*[c[mma_i_height, mma_i_width, 4+i] for i in range(4)])
out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
c_i = [c[mma_i_height, mma_i_width, i].store(out1.gep(i)) for i in range(4)] + [c[mma_i_height, mma_i_width, 4+i].store(out2.gep(i)) for i in range(4)]
c_store = UOp.group(*c_i).end(mma_i_height, mma_i_width, mma_i_inner)
out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
c_i = [c[height, width, i].store(out1.gep(i)) for i in range(4)] + [c[height, width, 4+i].store(out2.gep(i)) for i in range(4)]
c_store = UOp.group(*c_i).end(height, width, inner)
self.ker.push_store(c_store, c)
return c.after(c_store).reshape(c.shape) if after else c_store
def mma_ABt(self, c:UOp, a:UOp, b:UOp, after=True):
def mma_ABt(self, c:UOp|RT, a:UOp|RT, b:UOp|RT, after=True):
c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
assert self.warps == 1
mma_i_height = UOp.range(c.shape[-3], Group.mma_rid)
mma_i_width = UOp.range(c.shape[-2], Group.mma_rid+1)
mma_i_inner = UOp.range(a.shape[-2], Group.mma_rid+2, AxisType.REDUCE)
Group.mma_rid += 3
for height in self.ker.range(c.shape[-3], track=False):
for width in self.ker.range(c.shape[-2], track=False):
for inner in self.ker.range(a.shape[-2], AxisType.REDUCE, track=False):
wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
b_in1 = UOp.vectorize(*([b[width, inner, i] for i in range(2)] + [b[width, inner, 4+i] for i in range(2)]))
c_out1 = UOp.vectorize(*[c[height, width, i] for i in range(4)])
b_in2 = UOp.vectorize(*([b[width, inner, 2+i] for i in range(2)] + [b[width, inner, 6+i] for i in range(2)]))
c_out2 = UOp.vectorize(*[c[height, width, 4+i] for i in range(4)])
a_in = UOp.vectorize(*[a[mma_i_height, mma_i_inner, i] for i in range(8)])
b_in1 = UOp.vectorize(*([b[mma_i_width, mma_i_inner, i] for i in range(2)] + [b[mma_i_width, mma_i_inner, 4+i] for i in range(2)]))
c_out1 = UOp.vectorize(*[c[mma_i_height, mma_i_width, i] for i in range(4)])
b_in2 = UOp.vectorize(*([b[mma_i_width, mma_i_inner, 2+i] for i in range(2)] + [b[mma_i_width, mma_i_inner, 6+i] for i in range(2)]))
c_out2 = UOp.vectorize(*[c[mma_i_height, mma_i_width, 4+i] for i in range(4)])
out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
c_i = [c[mma_i_height, mma_i_width, i].store(out1.gep(i)) for i in range(4)] + [c[mma_i_height, mma_i_width, 4+i].store(out2.gep(i)) for i in range(4)]
c_store = UOp.group(*c_i).end(mma_i_height, mma_i_width, mma_i_inner)
out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
c_i = [c[height, width, i].store(out1.gep(i)) for i in range(4)] + [c[height, width, 4+i].store(out2.gep(i)) for i in range(4)]
c_store = UOp.group(*c_i).end(height, width, inner)
self.ker.push_store(c_store, c)
return c.after(c_store).reshape(c.shape) if after else c_store
map_rid = 400
def map(self, a:UOp, op:Callable[[UOp], UOp]|Callable[[UOp, tuple], UOp]):
def map(self, a:ALL_TILES, op:Callable[[UOp], UOp]|Callable[[UOp, tuple], UOp]):
a = cast(UOp, a)
assert self.warps == 1
rngs_for_shape = tuple(UOp.range(dim, Group.map_rid + i) for i, dim in enumerate(a.shape))
@@ -123,14 +120,12 @@ class Group:
self.ker.push_store(a_store, a)
return a.after(a_store).reshape(a.shape)
def row_reduce(self, vec:UOp, src:UOp, op:Callable[[UOp, UOp], UOp]):
def row_reduce(self, vec:UOp|RV, src:UOp|RT, op:Callable[[UOp, UOp], UOp]):
vec, src = cast(UOp, vec), cast(UOp, src)
assert self.warps == 1
red_local = UOp.placeholder((self.group_threads, 2), src.dtype.base, addrspace=AddrSpace.LOCAL, slot=slots.shared_slot)
slots.shared_slot += 1
red_reg = UOp.placeholder((2,), src.dtype.base, addrspace=AddrSpace.REG, slot=slots.register_slot)
slots.register_slot += 1
red_local = self.ker.alloc((self.group_threads, 2), src.dtype.base, AddrSpace.LOCAL)
red_reg = self.ker.alloc((2,), src.dtype.base, AddrSpace.REG)
for height in self.ker.range(src.shape[-3], track=False):
i = UOp.range(red_reg.size, Group.clear_rid)
@@ -139,28 +134,28 @@ class Group:
reg_store = red_reg.flatten()[i].store(0.).end(i)
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
for i_outer in self.ker.range(2, track=False):
for outer in self.ker.range(2, track=False):
for width in self.ker.range(src.shape[-2], AxisType.REDUCE, track=False):
for i_inner in self.ker.range(4, AxisType.REDUCE, track=False):
elem_index = i_inner + 2 * (i_inner // 2) + i_outer * 2
reg_store = red_reg[i_outer].store(op(red_reg[i_outer], src[height, width, elem_index])).end(i_inner, width, i_outer)
for inner in self.ker.range(4, AxisType.REDUCE, track=False):
elem_index = inner + 2 * (inner // 2) + outer * 2
reg_store = red_reg[outer].store(op(red_reg[outer], src[height, width, elem_index])).end(inner, width, outer)
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
# store to shared memory
for i_outer in self.ker.range(2, track=False):
red_local_store = red_local[self.laneid, i_outer].store(red_reg[i_outer]).end(i_outer)
for outer in self.ker.range(2, track=False):
red_local_store = red_local[self.laneid, outer].store(red_reg[outer]).end(outer)
red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
# reduce from shared memory
for i_outer in self.ker.range(2, track=False):
for i_inner in self.ker.range(3, AxisType.REDUCE, track=False):
offset = (self.laneid // 4) * 4 + ((self.laneid + i_inner + 1) % 4)
reg_store = red_reg[i_outer].store(op(red_reg[i_outer], red_local[offset, i_outer])).end(i_inner, i_outer)
for outer in self.ker.range(2, track=False):
for inner in self.ker.range(3, AxisType.REDUCE, track=False):
offset = (self.laneid // 4) * 4 + ((self.laneid + inner + 1) % 4)
reg_store = red_reg[outer].store(op(red_reg[outer], red_local[offset, outer])).end(inner, outer)
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
# reduce with vec
for i_outer in self.ker.range(2, track=False):
vec_store = vec[height, 0, i_outer].store(op(vec[height, 0, i_outer], red_reg[i_outer])).end(i_outer, height)
for outer in self.ker.range(2, track=False):
vec_store = vec[height, 0, outer].store(op(vec[height, 0, outer], red_reg[outer])).end(outer, height)
self.ker.push_store(vec_store, vec)
return vec.after(vec_store).reshape(vec.shape)
@@ -168,39 +163,37 @@ class Group:
# ops that can work across multiple warps
LOAD_INNER = 8
load_rid = 100
def load(self, dst:UOp, src:UOp, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
def load(self, dst:ALL_TILES, src:ALL_TILES, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
dst, src = cast(UOp, dst), cast(UOp, src)
assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
if dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.LOCAL:
srcf = src.flatten(-2)
load_i_height = UOp.range(dst.shape[-3], Group.load_rid)
load_i_width = UOp.range(dst.shape[-2], Group.load_rid+1)
load_i_inner = UOp.range(RT_BASE_TILE_NEPT, Group.load_rid+2)
Group.load_rid += 3
if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
else: local_warpid = self.warpid
warp_laneid = self.threadIdx_x % WARP_THREADS
if not transpose:
row = (local_warpid * dst.shape[-3] + load_i_height) * TILE_ROW_DIM + (warp_laneid // 4)
col = load_i_width * TILE_COL_DIM + 2 * (warp_laneid % 4)
for height in self.ker.range(dst.shape[-3], track=False):
for width in self.ker.range(dst.shape[-2], track=False):
for inner in self.ker.range(RT.BASE_TILE_NEPT, track=False):
if not transpose:
row = (local_warpid * dst.shape[-3] + height) * RT.BASE_TILE_ROWS + (warp_laneid // 4)
col = width * RT.BASE_TILE_COLS + 2 * (warp_laneid % 4)
row_offset = ((load_i_inner % 4) // 2) * 8
col_offset = (load_i_inner % 2) + (load_i_inner // 4) * 8
else:
row = (local_warpid * dst.shape[-3] + load_i_height) * TILE_ROW_DIM + 2 * (warp_laneid % 4)
col = load_i_width * TILE_COL_DIM + (warp_laneid // 4)
row_offset = ((inner % 4) // 2) * 8
col_offset = (inner % 2) + (inner // 4) * 8
else:
row = (local_warpid * dst.shape[-3] + height) * RT.BASE_TILE_ROWS + 2 * (warp_laneid % 4)
col = width * RT.BASE_TILE_COLS + (warp_laneid // 4)
row_offset = (load_i_inner % 2) + (load_i_inner // 4) * 8
col_offset = ((load_i_inner % 4) // 2) * 8
row_offset = (inner % 2) + (inner // 4) * 8
col_offset = ((inner % 4) // 2) * 8
src_i_last = (row + row_offset) * src.shape[-1] + col + col_offset
src_i_last = (row + row_offset) * src.shape[-1] + col + col_offset
dst_store = dst[*dst_idxs, load_i_height, load_i_width, load_i_inner].store(srcf[*idxs[:-2], src_i_last])
dst_store = dst_store.end(load_i_height, load_i_width, load_i_inner)
dst_store = dst[*dst_idxs, height, width, inner].store(srcf[*idxs[:-2], src_i_last])
dst_store = dst_store.end(height, width, inner)
elif dst_dtype.addrspace == AddrSpace.LOCAL and src_dtype.addrspace == AddrSpace.GLOBAL:
dstf = dst.flatten(-2)
@@ -214,50 +207,46 @@ class Group:
memcpy_per_row = dst.shape[-1] // Group.LOAD_INNER
total_calls = prod(dst.shape[-2:]) // (self.group_threads * Group.LOAD_INNER)
load_i_outer = UOp.range(total_calls, Group.load_rid)
load_i_inner = UOp.range(Group.LOAD_INNER, Group.load_rid+1)
Group.load_rid += 2
for outer in self.ker.range(total_calls, track=False):
for inner in self.ker.range(Group.LOAD_INNER, track=False):
load_idx = outer * self.group_threads + self.laneid
row = load_idx // memcpy_per_row
col = (load_idx * Group.LOAD_INNER) % dst.shape[-1]
load_idx = load_i_outer * self.group_threads + self.laneid
row = load_idx // memcpy_per_row
col = (load_idx * Group.LOAD_INNER) % dst.shape[-1]
dst_i = row * dst.shape[-1] + col + inner
src_i += row * row_stride + col + inner
dst_i = row * dst.shape[-1] + col + load_i_inner
src_i += row * row_stride + col + load_i_inner
dst_store = dstf[*dst_idxs, dst_i].store(srcf[src_i]).end(load_i_outer, load_i_inner)
dst_store = dstf[*dst_idxs, dst_i].store(srcf[src_i]).end(outer, inner)
else:
raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
return dst.after(dst_store.barrier()).reshape(dst.shape)
STORE_INNER = 8
store_rid = 200
def store(self, dst:UOp, src:UOp, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis=0, after=True):
def store(self, dst:ALL_TILES, src:ALL_TILES, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis=0, after=True):
dst, src = cast(UOp, dst), cast(UOp, src)
assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.LOCAL:
dstf = dst.flatten(-2)
store_i_height = UOp.range(src.shape[-3], Group.store_rid)
store_i_width = UOp.range(src.shape[-2], Group.store_rid+1)
store_i_inner = UOp.range(RT_BASE_TILE_NEPT, Group.store_rid+2)
Group.store_rid += 3
if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
else: local_warpid = self.warpid
warp_laneid = self.threadIdx_x % WARP_THREADS
row = (local_warpid * src.shape[-3] + store_i_height) * TILE_ROW_DIM + (warp_laneid // 4)
col = store_i_width * TILE_COL_DIM + 2 * (warp_laneid % 4)
for height in self.ker.range(src.shape[-3], track=False):
for width in self.ker.range(src.shape[-2], track=False):
for inner in self.ker.range(RT.BASE_TILE_NEPT, track=False):
row = (local_warpid * src.shape[-3] + height) * RT.BASE_TILE_ROWS + (warp_laneid // 4)
col = width * RT.BASE_TILE_COLS + 2 * (warp_laneid % 4)
row_offset = ((store_i_inner % 4) // 2) * 8
col_offset = (store_i_inner % 2) + (store_i_inner // 4) * 8
row_offset = ((inner % 4) // 2) * 8
col_offset = (inner % 2) + (inner // 4) * 8
dst_i_last = (row + row_offset) * dst.shape[-1] + col + col_offset
dst_i_last = (row + row_offset) * dst.shape[-1] + col + col_offset
dst_store = dstf[*idxs[:-2], dst_i_last].store(src[*src_idxs, store_i_height, store_i_width, store_i_inner])
dst_store = dst_store.end(store_i_height, store_i_width, store_i_inner)
dst_store = dstf[*idxs[:-2], dst_i_last].store(src[*src_idxs, height, width, inner])
dst_store = dst_store.end(height, width, inner)
elif src_dtype.addrspace == AddrSpace.LOCAL and dst_dtype.addrspace == AddrSpace.GLOBAL:
dstf = dst.flatten()
row_stride = prod(dst.shape[axis+1:])
@@ -271,18 +260,16 @@ class Group:
memcpy_per_row = src.shape[-1] // Group.STORE_INNER
total_calls = prod(src.shape[-2:]) // (self.group_threads * Group.STORE_INNER)
store_i_outer = UOp.range(total_calls, Group.store_rid)
store_i_inner = UOp.range(Group.STORE_INNER, Group.store_rid+1)
Group.store_rid += 2
for outer in self.ker.range(total_calls, track=False):
for inner in self.ker.range(Group.STORE_INNER, track=False):
load_idx = outer * self.group_threads + self.laneid
row = load_idx // memcpy_per_row
col = (load_idx * Group.STORE_INNER) % src.shape[-1]
load_idx = store_i_outer * self.group_threads + self.laneid
row = load_idx // memcpy_per_row
col = (load_idx * Group.STORE_INNER) % src.shape[-1]
src_i = row * src.shape[-1] + col + inner
dst_i += row * row_stride + col + inner
src_i = row * src.shape[-1] + col + store_i_inner
dst_i += row * row_stride + col + store_i_inner
dst_store = dstf[dst_i].store(srcf[*src_idxs, src_i]).end(store_i_outer, store_i_inner)
dst_store = dstf[dst_i].store(srcf[*src_idxs, src_i]).end(outer, inner)
else:
raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
extra/thunder/tiny/tk/kernel.py
+33 -2
View File
@@ -1,7 +1,8 @@
from contextlib import AbstractContextManager
from tinygrad.uop.ops import UOp, KernelInfo, AxisType
from tinygrad.uop.ops import UOp, KernelInfo, AxisType, AddrSpace
from extra.thunder.tiny.tk import WARP_THREADS
from extra.thunder.tiny.tk.group import Group
from extra.thunder.tiny.tk.tiles import GL, ST, RT, RV
class _tk_range:
user_rid = 0
@@ -25,6 +26,11 @@ class Kernel(AbstractContextManager):
self.range_stack = []
self.store_stack = []
self.global_slot = 0
self.shared_slot = 0
self.register_slot = 0
self.allocs = {}
@property
def warpid(self): return self.threadIdx_x // WARP_THREADS
@@ -42,6 +48,31 @@ class Kernel(AbstractContextManager):
if track: self.range_stack.append(rng)
return rng
def alloc(self, shape, dtype, addrspace:AddrSpace, name:str|None=None):
match addrspace:
case AddrSpace.GLOBAL:
slot = self.global_slot
self.global_slot += 1
case AddrSpace.LOCAL:
slot = self.shared_slot
self.shared_slot += 1
case AddrSpace.REG:
slot = self.register_slot
self.register_slot += 1
uop = UOp.placeholder(shape, dtype, slot=slot, addrspace=addrspace)
if name:
if (name, shape) in self.allocs: return self.allocs[(name, shape)]
self.allocs[(name, shape)] = uop
return uop
def gl(self, shape, dtype): return GL.create(shape, dtype, self)
def st(self, shape, dtype): return ST.create(shape, dtype, self)
def rt(self, shape, dtype): return RT.create(shape, dtype, self)
def rv(self, length, dtype, layout="naive"): return RV.create(length, dtype, layout, self)
def push_store(self, store:UOp, uop:UOp): self.store_stack.append((store, uop))
def finish(self):
@@ -49,7 +80,7 @@ class Kernel(AbstractContextManager):
rngs = []
while self.range_stack: rngs.append(self.range_stack.pop(0)._rng)
return self.store_stack.pop()[0].end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
return self.store_stack.pop()[0]._uop.end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
def endrange(self):
last_store = self.store_stack.pop()
extra/thunder/tiny/tk/tiles.py
+133 -42
View File
@@ -1,52 +1,143 @@
import math
from typing import cast, Callable
from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
from tinygrad.uop.ops import AxisType, UOp, KernelInfo, Ops
from tinygrad.engine.realize import ExecItem, get_runner
from tinygrad.dtype import AddrSpace, PtrDType
from tinygrad.helpers import getenv, prod
import functools
from tinygrad.dtype import AddrSpace
from tinygrad.mixin import MathMixin
from tinygrad.uop.ops import UOp, Ops
from extra.thunder.tiny.tk import WARP_THREADS
class _Slots:
def __init__(self):
self.global_slot = 0
self.shared_slot = 0
self.register_slot = 0
slots = _Slots()
def unwrap(x):
if hasattr(x, "_uop"): return x._uop
if isinstance(x, (list, tuple)): return type(x)(unwrap(y) for y in x)
if isinstance(x, dict): return {k: unwrap(v) for k,v in x.items()}
return x
def gl(shape, dtype):
slots.global_slot += 1
return UOp.placeholder(shape, dtype, slot=slots.global_slot-1)
def wrap(x, ker, cls):
if isinstance(x, UOp): return cls(x, ker)
if isinstance(x, (list, tuple)): return type(x)(wrap(y, ker, cls) for y in x)
return x
shared_slot = 0
def st(shape, dtype):
slots.shared_slot += 1
return UOp.placeholder(shape, dtype, addrspace=AddrSpace.LOCAL, slot=slots.shared_slot-1)
def autowrap(source_cls, blacklist=None):
if blacklist is None:
blacklist = {
"__init__", "__new__", "__str__", "__del__", "__repr__", "__dict__", "__getattribute__",
"__setattr__", "__delattr__", "__weakref__", "__slots__", "__class__",
"__reduce__", "__reduce_ex__", "__getstate__", "__setstate__", "__hash__"
}
TILE_ROW_DIM, TILE_COL_DIM = 16, 16
RT_BASE_TILE_NE = TILE_ROW_DIM * TILE_COL_DIM
RT_BASE_TILE_NEPT = RT_BASE_TILE_NE // WARP_THREADS
register_slot = 0
def rt(shape, dtype):
assert len(shape) == 2
def decorator(cls):
def __getattr__(self, name):
uop = object.__getattribute__(self, "_uop")
val = getattr(uop, name)
if callable(val):
@functools.wraps(val)
def proxy(*args, **kwargs):
return wrap(val(*unwrap(args), **unwrap(kwargs)), self.ker, cls)
return proxy
if name in UOp.__slots__: return val
return wrap(val, self.ker, cls)
cls.__getattr__ = __getattr__
height = shape[0] // TILE_ROW_DIM
width = shape[1] // TILE_COL_DIM
for name in dir(source_cls):
if name in blacklist or not name.startswith("__"): continue
slots.register_slot += 1
return UOp.placeholder((height, width, RT_BASE_TILE_NEPT), dtype, addrspace=AddrSpace.REG, slot=slots.register_slot-1)
for base in cls.mro():
if base is source_cls: break
if name in base.__dict__: break
else:
original = getattr(source_cls, name)
if callable(original):
def make_proxy(op_name, func):
def proxy(self, *args, **kwargs):
return wrap(func(self._uop, *unwrap(args), **unwrap(kwargs)), self.ker, cls)
return proxy
setattr(cls, name, make_proxy(name, original))
def rv(length, dtype, layout="naive"):
tiles = length // TILE_ROW_DIM
match layout:
case "naive":
inner_dim = 1
outer_dim = (tiles + 1) // 2
case "ortho":
inner_dim = 1
outer_dim = tiles
case _: raise NotImplementedError(f"rv layout {layout} not implemented")
return cls
return decorator
slots.register_slot += 1
return UOp.placeholder((outer_dim, inner_dim, 2), dtype, addrspace=AddrSpace.REG, slot=slots.register_slot-1)
class TileMathMixin(MathMixin):
def alu(self, op, *src, inner_op=lambda x:x):
assert isinstance(self, (RT, RV))
if len(src) == 0:
if self._uop._shape is None: uop = UOp.alu(self._uop, op)
else: uop = self.ker.warp.map(self._uop, lambda x: UOp.alu(x, op))
elif len(src) == 1:
if self._uop._shape is None: uop = UOp.alu(self._uop, op, inner_op(self._uop.ufix(src[0])))
elif isinstance(src[0], (int,float,bool)): uop = self.ker.warp.map(self._uop, lambda x: UOp.alu(x, op, inner_op(x.ufix(src[0]))))
elif src[0]._shape is None: uop = UOp.alu(self._uop, op, inner_op(self._uop.ufix(src[0])))
else:
if isinstance(self, RT) and isinstance(src[0], RV): uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[idx[0], 0, (idx[2]%4)//2])))
else: uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[*idx])))
else: raise NotImplementedError
return type(self)(uop, self.ker)
def const_like(self, b): return b
# override ops that do compute on the src uop
def sub(self, x, reverse=False):
return self.ufix(x).alu(Ops.ADD, self, inner_op=lambda y: -y) if reverse else self.alu(Ops.ADD, self.ufix(x), inner_op=lambda y: -y)
def div(self, x, reverse=False):
return self.ufix(x).alu(Ops.MUL, self, inner_op=lambda y: 1/y) if reverse else self.alu(Ops.MUL, self.ufix(x), inner_op=lambda y: 1/y)
@autowrap(UOp)
class GL:
def __init__(self, uop, ker):
self._uop, self.ker = uop, ker
@classmethod
def create(cls, shape, dtype, ker):
uop = ker.alloc(shape, dtype, AddrSpace.GLOBAL)
return cls(uop, ker)
@autowrap(UOp)
class ST:
def __init__(self, uop, ker):
self._uop, self.ker = uop, ker
@classmethod
def create(cls, shape, dtype, ker):
uop = ker.alloc(shape, dtype, AddrSpace.LOCAL)
return cls(uop, ker)
@autowrap(UOp)
class RT(TileMathMixin):
BASE_TILE_ROWS, BASE_TILE_COLS = 16, 16
BASE_TILE_NE = BASE_TILE_ROWS * BASE_TILE_COLS
BASE_TILE_NEPT = BASE_TILE_NE // WARP_THREADS
def __init__(self, uop, ker):
self._uop, self.ker = uop, ker
@classmethod
def create(cls, shape, dtype, ker):
assert len(shape) == 2
assert shape[0] % RT.BASE_TILE_ROWS == 0
assert shape[1] % RT.BASE_TILE_COLS == 0
height = shape[0] // RT.BASE_TILE_ROWS
width = shape[1] // RT.BASE_TILE_COLS
uop = ker.alloc((height, width, RT.BASE_TILE_NEPT), dtype, AddrSpace.REG)
return cls(uop, ker)
@autowrap(UOp)
class RV(TileMathMixin):
def __init__(self, uop, ker):
self._uop, self.ker = uop, ker
@classmethod
def create(cls, length, dtype, layout, ker):
tiles = length // RT.BASE_TILE_ROWS
match layout:
case "naive":
inner_dim = 1
outer_dim = (tiles + 1) // 2
case "ortho":
inner_dim = 1
outer_dim = tiles
case _: raise NotImplementedError(f"rv layout {layout} not implemented")
uop = ker.alloc((outer_dim, inner_dim, 2), dtype, AddrSpace.REG)
return RV(uop, ker)
ALL_TILES = UOp | GL | ST | RT | RV
extra/usbgpu/patch.py
-1
View File
@@ -1,6 +1,5 @@
#!/usr/bin/env python3
import sys, os, zlib, struct, hashlib
from hexdump import hexdump
from tinygrad.helpers import DEBUG, getenv, fetch
from tinygrad.runtime.support.usb import USB3
pyproject.toml
+1 -2
View File
@@ -30,7 +30,6 @@ packages = [
'tinygrad.runtime',
'tinygrad.runtime.autogen',
'tinygrad.runtime.autogen.am',
'tinygrad.runtime.autogen.nv',
'tinygrad.runtime.graph',
'tinygrad.runtime.support',
'tinygrad.runtime.support.am',
@@ -162,7 +161,6 @@ exclude = [
".git/",
"docs/",
"extra/",
"tinygrad/runtime/autogen",
"test/external/mlperf_resnet",
"test/external/mlperf_unet3d",
]
@@ -228,6 +226,7 @@ select = [
"F541",
"F841",
]
"tinygrad/runtime/autogen/**/*.py" = ["E501", "F401", "E722", "E731", "F821", "A006"]
[tool.ruff.format]
exclude = ["*"]
test/external/external_test_onnx_backend.py
Vendored
-1
View File
@@ -112,7 +112,6 @@ backend_test.exclude('test_dequantizelinear_e5m2_cpu')
backend_test.exclude('test_dequantizelinear_float4e2m1_cpu')
# we don't support indexes
backend_test.exclude('test_nonzero_*')
# no support for int pow
backend_test.exclude('test_pow_types_int32_int32_cpu')
test/mockgpu/amd/amdgpu.py
+2 -2
View File
@@ -33,7 +33,7 @@ remu = _try_dlopen_remu()
def create_sdma_packets():
# TODO: clean up this, if we want to keep it
structs = {}
for name,pkt in [(name,s) for name,s in amd_gpu.__dict__.items() if name.startswith("struct_SDMA_PKT_") and name.endswith("_TAG")]:
for name,pkt in [(name,s) for name,s in amd_gpu.__dict__.items() if name.startswith("rocr_AMD_SDMA_PKT_") and name.endswith("_TAG")]:
names = set()
fields = []
for pkt_fields in pkt._fields_:
@@ -47,7 +47,7 @@ def create_sdma_packets():
# merge together 64-bit fields, otherwise just append them
if fname.endswith("_63_32") and fields[-1][0].endswith("_31_0"): fields[-1] = tuple([fname[:-6], ctypes.c_ulong, 64])
else: fields.append(tuple([fname, *union_fields[1:]]))
new_name = name[16:-4].lower()
new_name = name[18:-4].lower()
structs[new_name] = init_c_struct_t(tuple(fields))
assert ctypes.sizeof(structs[new_name]) == ctypes.sizeof(pkt), f"{ctypes.sizeof(structs[new_name])} != {ctypes.sizeof(pkt)}"
return type("SDMA_PKTS", (object, ), structs)
test/mockgpu/cuda/cuda.py
+1 -1
View File
@@ -164,7 +164,7 @@ def cuStreamWaitEvent(stream: Any, event, flags: int) -> int: return orig_cuda.C
def cuCtxSynchronize() -> int: return orig_cuda.CUDA_SUCCESS
def cuGetErrorString(error: int, pStr) -> int:
error_str = orig_cuda.cudaError_enum__enumvalues.get(error, "Unknown CUDA error").encode()
error_str = orig_cuda.enum_cudaError_enum.get(error, "Unknown CUDA error").encode()
buf = ctypes.create_string_buffer(error_str)
# Set the pointer to point to our error string buffer
pStr._obj.value = ctypes.cast(buf, ctypes.POINTER(ctypes.c_char))
test/mockgpu/nv/nvdriver.py
+9 -4
View File
@@ -1,5 +1,5 @@
import ctypes, mmap, collections, functools, os
import tinygrad.runtime.autogen.nv_gpu as nv_gpu
from tinygrad.runtime.autogen import nv_570 as nv_gpu
from typing import Any
from tinygrad.helpers import to_mv
from test.mockgpu.driver import VirtDriver, VirtFileDesc, VirtFile
@@ -153,8 +153,10 @@ class NVDriver(VirtDriver):
51059, 51069, 51071, 51632, 51639, 51639, 51706, 52019, 222, 50287, 50273, 50031, 50017] # from ada102
params.numClasses = len(classes)
if struct.cmd == nv_gpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST:
clslist = to_mv(params.classList, params.numClasses * 4).cast('I')
for i,c in enumerate(classes): clslist[i] = c
if params.classList and params.numClasses > 0:
clslist = to_mv(params.classList, params.numClasses * 4).cast('I')
for i,c in enumerate(classes): clslist[i] = c
else: params.numClasses = len(classes)
else:
for i,c in enumerate(classes): params.classList[i] = c
elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_GR_GET_INFO:
@@ -192,6 +194,9 @@ class NVDriver(VirtDriver):
params.mmuFaultInfoList[0].faultAddress = int(os.environ['MOCKGPU_EMU_FAULTADDR'], base=16)
params.mmuFaultInfoList[0].faultType = 1
params.mmuFaultInfoList[0].accessType = 1
elif struct.cmd == nv_gpu.NV0000_CTRL_CMD_SYSTEM_GET_BUILD_VERSION_V2:
params = nv_gpu.NV0000_CTRL_SYSTEM_GET_BUILD_VERSION_V2_PARAMS.from_address(params_ptr)
params.driverVersionBuffer = b"570.00.00\0"
else: raise RuntimeError(f"Unknown {struct.cmd} to rm_control")
return 0
@@ -254,4 +259,4 @@ class NVDriver(VirtDriver):
for gpu in self.gpus.values():
for q in gpu.queues:
if q.ctrl.GPGet != q.ctrl.GPPut:
any_progress |= q.execute()
any_progress |= q.execute()
test/mockgpu/nv/nvgpu.py
+1 -1
View File
@@ -1,5 +1,5 @@
import ctypes, time
import tinygrad.runtime.autogen.nv_gpu as nv_gpu
from tinygrad.runtime.autogen import nv_570 as nv_gpu
from enum import Enum, auto
from test.mockgpu.gpu import VirtGPU
from test.mockgpu.helpers import _try_dlopen_gpuocelot
test/test_linearizer.py
+1
View File
@@ -78,6 +78,7 @@ class TestLinearizer(unittest.TestCase):
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
assert len(ranges) == 1 # NOTE: it collapses now
@unittest.expectedFailure # TODO: investigate
def test_two_nested_range_alt_indexing(self):
a = Tensor([2, 2]).realize()
out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
test/test_tensor.py
+11 -1
View File
@@ -3,7 +3,7 @@ import torch
import unittest, copy, mmap, random, math, array
from tinygrad import Tensor, Device, dtypes
from tinygrad.tensor import _METADATA
from tinygrad.helpers import getenv, temp, mv_address
from tinygrad.helpers import Context, getenv, temp, mv_address
from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
from hypothesis import given, settings, strategies as strat
from tinygrad.device import is_dtype_supported
@@ -846,6 +846,16 @@ class TestTensorMetadata(unittest.TestCase):
#self.assertEqual(len(bw), 1)
#self.assertEqual(bw[0].name, "sigmoid")
def test_tracemeta_0(self):
with Context(TRACEMETA=0):
x = Tensor.rand(3, requires_grad=True)
y = Tensor.rand(3, requires_grad=True)
out = (x.relu() * y.sigmoid()).sum()
self.assertIsNone(out.uop.metadata)
self.assertIsNone(out.uop.src[0].metadata)
si = out.schedule()[-1]
self.assertEqual(si.metadata, ())
class TestIdxUpcast(unittest.TestCase):
def _find_op(self, ast: UOp, op: Ops):
if ast.op is op: return ast
test/test_tiny.py
+4 -6
View File
@@ -32,8 +32,8 @@ class TestTiny(unittest.TestCase):
self.assertListEqual(out.tolist(), [2]*16)
def test_cat(self):
out = Tensor.cat(Tensor.ones(8).contiguous(), Tensor.ones(8).contiguous())
self.assertListEqual(out.tolist(), [1]*16)
out = Tensor.cat(Tensor.ones(8).contiguous(), Tensor.zeros(8).contiguous())
self.assertListEqual(out.tolist(), [1]*8+[0]*8)
def test_sum(self):
out = Tensor.ones(256).contiguous().sum()
@@ -62,7 +62,7 @@ class TestTiny(unittest.TestCase):
out = Tensor.rand(10)
for x in out.tolist():
self.assertGreaterEqual(x, 0.0)
self.assertLessEqual(x, 1.0)
self.assertLess(x, 1.0)
# *** JIT (for Python speed) ***
@@ -138,9 +138,7 @@ class TestTiny(unittest.TestCase):
nn.Conv2d(8, 8, 5), Tensor.relu]
# replace random weights with ones
# TODO: there's a bug here where it's tying two of the biases together. we need UNIQUE const
#Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
for p in nn.state.get_parameters(layers): p.replace(Tensor.empty(p.shape))
Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
# realize gradients
for x in nn.state.get_parameters(layers): x.requires_grad_()
test/external/external_test_tk.py → test/testextra/test_tk.py
+111 -69
View File
@@ -2,30 +2,34 @@ import unittest, math
from tinygrad import Tensor, Device, dtypes, Context
from tinygrad.engine.realize import ExecItem, get_runner
from tinygrad.helpers import CI
from tinygrad.renderer.ptx import PTXRenderer
import numpy as np
from extra.thunder.tiny.tk import WARP_THREADS
from extra.thunder.tiny.tk.kernel import Kernel
from extra.thunder.tiny.tk.tiles import gl, st, rt, rv
@unittest.skipUnless(Device.DEFAULT in ["CUDA", "NV"], "only cuda")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "no ptx")
class TestTK(unittest.TestCase):
@unittest.skipIf(CI, "no wmma in ci")
def test_simple_matmul(self):
N = 32
BLOCK_SIZE = 16
with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
warp = ker.warp
c = gl((1, 1, N, N), dtypes.float32)
a = gl((1, 1, N, N), dtypes.bfloat16)
b = gl((1, 1, N, N), dtypes.bfloat16)
c = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.bfloat16)
b = ker.gl((1, 1, N, N), dtypes.bfloat16)
a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
c_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
c_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
c_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
col, row = ker.blockIdx_x, ker.blockIdx_y
@@ -59,23 +63,24 @@ class TestTK(unittest.TestCase):
np.testing.assert_allclose(c.numpy(), ref.numpy())
@unittest.skipIf(CI, "no wmma in ci")
def test_simple_matmul_transposed(self):
N = 32
BLOCK_SIZE = 16
with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
warp = ker.warp
c = gl((1, 1, N, N), dtypes.float32)
a = gl((1, 1, N, N), dtypes.bfloat16)
b = gl((1, 1, N, N), dtypes.bfloat16)
c = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.bfloat16)
b = ker.gl((1, 1, N, N), dtypes.bfloat16)
a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
c_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
c_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
c_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
col, row = ker.blockIdx_x, ker.blockIdx_y
@@ -115,14 +120,14 @@ class TestTK(unittest.TestCase):
with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
warp = ker.warp
b = gl((1, 1, N, N), dtypes.float32)
a = gl((1, 1, N, N), dtypes.float32)
b = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
col, row = ker.blockIdx_x, ker.blockIdx_y
@@ -147,22 +152,60 @@ class TestTK(unittest.TestCase):
np.testing.assert_allclose(b.numpy(), ref.numpy())
def test_add(self):
N = 32
BLOCK_SIZE = 16
with Kernel((1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
for tile_row in ker.range(N // BLOCK_SIZE):
for tile_col in ker.range(N // BLOCK_SIZE):
a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
a_reg = warp.load(a_reg, a_smem)
a_reg += 1
a_smem = warp.store(a_smem, a_reg)
b = warp.store(b, a_smem, (0, 0, tile_row, tile_col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
for _ in range(5): ei.run(wait=True)
b = b.float()
ref = a.float() + 1
np.testing.assert_allclose(b.numpy(), ref.numpy())
def test_max(self):
N = 16
BLOCK_SIZE = 16
with Kernel((1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = gl((1, 1, N, N), dtypes.float32)
a = gl((1, 1, N, N), dtypes.float32)
b = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
max_reg = rv(BLOCK_SIZE, dtypes.float32, "ortho")
max_reg = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
for tile_row in ker.range(N // BLOCK_SIZE):
max_reg = warp.neg_inf(max_reg.after(tile_row))
@@ -200,16 +243,16 @@ class TestTK(unittest.TestCase):
with Kernel((1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = gl((1, 1, N, M), dtypes.float32)
a = gl((1, 1, N, M), dtypes.float32)
b = ker.gl((1, 1, N, M), dtypes.float32)
a = ker.gl((1, 1, N, M), dtypes.float32)
a_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
b_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
a_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
b_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
a_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
b_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
max_reg = rv(BLOCK_N, dtypes.float32, "ortho")
max_reg = ker.rv(BLOCK_N, dtypes.float32, "ortho")
for tile_row in ker.range(N // BLOCK_N):
max_reg = warp.neg_inf(max_reg.after(tile_row))
@@ -247,16 +290,16 @@ class TestTK(unittest.TestCase):
with Kernel((1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = gl((1, 1, N, N), dtypes.float32)
a = gl((1, 1, N, N), dtypes.float32)
b = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
sum_reg = rv(BLOCK_SIZE, dtypes.float32, "ortho")
sum_reg = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
for tile_row in ker.range(N // BLOCK_SIZE):
sum_reg = warp.zero(sum_reg.after(tile_row))
@@ -294,16 +337,16 @@ class TestTK(unittest.TestCase):
with Kernel((1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = gl((1, 1, N, M), dtypes.float32)
a = gl((1, 1, N, M), dtypes.float32)
b = ker.gl((1, 1, N, M), dtypes.float32)
a = ker.gl((1, 1, N, M), dtypes.float32)
a_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
b_smem = st((BLOCK_N, BLOCK_M), dtypes.float32)
a_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
b_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
a_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
b_reg = rt((BLOCK_N, BLOCK_M), dtypes.float32)
a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
sum_reg = rv(BLOCK_N, dtypes.float32, "ortho")
sum_reg = ker.rv(BLOCK_N, dtypes.float32, "ortho")
for tile_row in ker.range(N // BLOCK_N):
sum_reg = warp.zero(sum_reg.after(tile_row))
@@ -335,22 +378,23 @@ class TestTK(unittest.TestCase):
np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
@unittest.skip("fake range not ended")
def test_softmax(self):
N = 32
BLOCK_SIZE = 16
with Kernel((1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
a = gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
b = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
a = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
a_smem = st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
max_vec_last = rv(BLOCK_SIZE, dtypes.float32, "ortho")
max_vec = rv(BLOCK_SIZE, dtypes.float32, "ortho")
norm_vec = rv(BLOCK_SIZE, dtypes.float32, "ortho")
max_vec_last = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
max_vec = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
norm_vec = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
max_vec = warp.neg_inf(max_vec)
norm_vec = warp.zero(norm_vec)
@@ -359,13 +403,13 @@ class TestTK(unittest.TestCase):
a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
a_reg = warp.load(a_reg, a_smem)
a_reg = warp.map(a_reg, lambda x: x * (1.0 / math.log(2)))
a_reg *= 1.0 / math.log(2)
max_vec_last = warp.copy(max_vec_last.after(tile_col), max_vec)
max_vec = warp.row_reduce(max_vec, a_reg, lambda a, b: a.maximum(b))
a_reg = warp.map(a_reg, lambda x, idx: (x - max_vec[idx[0], 0, (idx[2]%4)//2]).exp2())
max_vec_last = warp.map(max_vec_last, lambda x, idx: (x - max_vec[*idx]).exp2())
norm_vec = warp.map(norm_vec, lambda x, idx: x * max_vec_last[*idx])
a_reg = (a_reg - max_vec).exp2()
max_vec_last = (max_vec_last - max_vec).exp2()
norm_vec *= max_vec_last
norm_vec = warp.row_reduce(norm_vec, a_reg, lambda a, b: a + b)
norm_vec = ker.endrange()
@@ -373,9 +417,9 @@ class TestTK(unittest.TestCase):
a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
a_reg = warp.load(a_reg, a_smem)
a_reg = warp.map(a_reg, lambda x: x * (1.0 / math.log(2)))
a_reg = warp.map(a_reg, lambda x, idx: (x - max_vec[idx[0], 0, (idx[2]%4)//2]).exp2())
a_reg = warp.map(a_reg, lambda x, idx: x / norm_vec[idx[0], 0, (idx[2]%4)//2])
a_reg *= 1.0 / math.log(2)
a_reg = (a_reg - max_vec).exp2()
a_reg /= norm_vec
a_smem = warp.store(a_smem, a_reg)
b = warp.store(b, a_smem, (0, 0, 0, tile_col), (), axis=2)
@@ -390,10 +434,8 @@ class TestTK(unittest.TestCase):
ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
for _ in range(5): ei.run(wait=True)
b = b.float()
print(b.tolist())
ref = a.float().softmax(axis=3)
print(ref.tolist())
np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
test/unit/test_autogen.py
+38
View File
@@ -44,4 +44,42 @@ class TestAutogen(unittest.TestCase):
test.argtypes = [Baz]
self.assertEqual(test(b), b.a + b.b + b.c + b.d)
@unittest.skipIf(WIN, "doesn't compile on windows")
def test_packed_structs(self):
NvU32 = ctypes.c_uint32
NvU64 = ctypes.c_uint64
class FWSECLIC_READ_VBIOS_DESC(Struct): pass
FWSECLIC_READ_VBIOS_DESC._packed_ = True
FWSECLIC_READ_VBIOS_DESC._fields_ = [
('version', NvU32),
('size', NvU32),
('gfwImageOffset', NvU64),
('gfwImageSize', NvU32),
('flags', NvU32),
]
class FWSECLIC_FRTS_REGION_DESC(Struct): pass
FWSECLIC_FRTS_REGION_DESC._packed_ = True
FWSECLIC_FRTS_REGION_DESC._fields_ = [
('version', NvU32),
('size', NvU32),
('frtsRegionOffset4K', NvU32),
('frtsRegionSize', NvU32),
('frtsRegionMediaType', NvU32),
]
class FWSECLIC_FRTS_CMD(Struct): pass
FWSECLIC_FRTS_CMD._packed_ = True
FWSECLIC_FRTS_CMD._fields_ = [
('readVbiosDesc', FWSECLIC_READ_VBIOS_DESC),
('frtsRegionDesc', FWSECLIC_FRTS_REGION_DESC),
]
read_vbios_desc = FWSECLIC_READ_VBIOS_DESC(version=0x1, size=ctypes.sizeof(FWSECLIC_READ_VBIOS_DESC), flags=2)
frst_reg_desc = FWSECLIC_FRTS_REGION_DESC(version=0x1, size=ctypes.sizeof(FWSECLIC_FRTS_REGION_DESC),
frtsRegionOffset4K=0xdead, frtsRegionSize=0x100, frtsRegionMediaType=2)
frts_cmd = FWSECLIC_FRTS_CMD(readVbiosDesc=read_vbios_desc, frtsRegionDesc=frst_reg_desc)
assert int.from_bytes(frts_cmd, 'little') == 0x2000001000000dead0000001400000001000000020000000000000000000000000000001800000001
assert int.from_bytes(frts_cmd.readVbiosDesc, 'little') == int.from_bytes(read_vbios_desc, 'little')
assert int.from_bytes(frts_cmd.frtsRegionDesc, 'little') == int.from_bytes(frst_reg_desc, 'little')
assert frts_cmd.readVbiosDesc.__class__ is FWSECLIC_READ_VBIOS_DESC
assert frts_cmd.frtsRegionDesc.__class__ is FWSECLIC_FRTS_REGION_DESC
if __name__ == "__main__": unittest.main()
test/unit/test_conv.py
+1
View File
@@ -62,6 +62,7 @@ class TestConv(unittest.TestCase):
np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0), atol=1e-5)
np.testing.assert_allclose(r2.numpy(), np.where(out.numpy() > 0, out.numpy(), (np.exp(out.numpy()) - 1)), atol=1e-5)
@unittest.skip("this test is flaky")
def test_two_overlapping_binops_no_rerun_wino(self):
with Context(WINO=1):
x = Tensor.randn(1,4,16,16)
test/unit/test_device.py
+3 -3
View File
@@ -81,20 +81,20 @@ class TestCompiler(unittest.TestCase):
def test_compile_cached(self):
diskcache_put("key", "123", None) # clear cache
getenv.cache_clear()
with Context(DISABLE_COMPILER_CACHE=0):
with Context(CCACHE=1):
self.assertEqual(MockCompiler("key").compile_cached("123"), str.encode("123"))
self.assertEqual(diskcache_get("key", "123"), str.encode("123"))
def test_compile_cached_disabled(self):
diskcache_put("disabled_key", "123", None) # clear cache
getenv.cache_clear()
with Context(DISABLE_COMPILER_CACHE=1):
with Context(CCACHE=0):
self.assertEqual(MockCompiler("disabled_key").compile_cached("123"), str.encode("123"))
self.assertIsNone(diskcache_get("disabled_key", "123"))
def test_device_compile(self):
getenv.cache_clear()
with Context(DISABLE_COMPILER_CACHE=1):
with Context(CCACHE=0):
a = Tensor([0.,1.], device=Device.DEFAULT).realize()
(a + 1).realize()
test/unit/test_viz.py
+6
View File
@@ -436,6 +436,12 @@ class TestVizProfiler(BaseTestViz):
sz = len(get_profile(prof))
self.assertLessEqual(sz/n_events, 26)
def test_calltrace(self):
def fxn(): return Tensor.empty(10).mul(2).realize()
fxn()
trace = get_viz_list()[0]["steps"][0]["trace"]
assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno == l for f,l,*_ in trace), str(trace)
# can pack up to 1hr 11 min of trace events
def test_trace_duration(self):
dur_mins = 72
tinygrad/device.py
+2 -2
View File
@@ -4,7 +4,7 @@ from collections import defaultdict
from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator
import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM
from tinygrad.helpers import Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM, select_first_inited
from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
from tinygrad.renderer import Renderer
@@ -266,7 +266,7 @@ class LRUAllocator(Allocator, Generic[DeviceType]):
class CompileError(Exception): pass
class Compiler:
def __init__(self, cachekey:str|None=None): self.cachekey = None if DISABLE_COMPILER_CACHE else cachekey
def __init__(self, cachekey:str|None=None): self.cachekey = cachekey if CCACHE else None
def compile(self, src:str) -> bytes: return src.encode() # NOTE: empty compiler is the default
def compile_cached(self, src:str) -> bytes:
if self.cachekey is None or (lib := diskcache_get(self.cachekey, src)) is None:
tinygrad/helpers.py
+22 -2
View File
@@ -173,13 +173,12 @@ TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS
SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
PICKLE_BUFFERS, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("LRU", 1)
CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
DISABLE_COMPILER_CACHE = ContextVar("DISABLE_COMPILER_CACHE", 0)
VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
EMULATE = ContextVar("EMULATE", "")
CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 1)
CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
VIZ = PROFILE = ContextVar("VIZ", 0)
SPEC = ContextVar("SPEC", 1)
# TODO: disable by default due to speed
@@ -188,6 +187,8 @@ PCONTIG = ContextVar("PCONTIG", 0) # partial contiguous in rangeify
DEBUG_RANGEIFY = ContextVar("DEBUG_RANGEIFY", 0)
# set to 1, this uses tuplize in the linearizer sort order
TUPLE_ORDER = ContextVar("TUPLE_ORDER", 1)
# set to 0 to disable the compiler cache
CCACHE = ContextVar("CCACHE", 1)
@dataclass(frozen=True)
class Metadata:
@@ -240,11 +241,29 @@ class Profiling(contextlib.ContextDecorator):
def perf_counter_us() -> decimal.Decimal: return decimal.Decimal(time.perf_counter_ns())/1000
@functools.cache
def lines(fn) -> list[str]:
try:
with open(fn, encoding="utf-8") as f: return f.readlines()
except (FileNotFoundError, OSError): return []
def printable(loc:tuple[str, int]) -> str:
try: return lines(loc[0])[loc[1]-1].strip()
except IndexError: return "<missing>"
def get_stacktrace(frm, max_frames=30) -> tuple[tuple, ...]:
ret:list[tuple] = []
for i in range(max_frames):
if (frm:=frm.f_back) is None: break
ret.append(((fc:=frm.f_code).co_filename, frm.f_lineno, fc.co_name, printable((fc.co_filename, frm.f_lineno))))
return tuple(ret)
@dataclass(frozen=True)
class TracingKey:
display_name:str # display name of this trace event
keys:tuple[Any, ...]=() # optional keys to search for related traces
ret:Any=None
tb:tuple[tuple, ...]|None=field(default_factory=lambda: get_stacktrace(sys._getframe(1)) if VIZ else None)
class ProfileEvent: pass
@@ -397,6 +416,7 @@ def to_mv(ptr:int, sz:int) -> memoryview: return memoryview((ctypes.c_uint8 * sz
def mv_address(mv): return ctypes.addressof(ctypes.c_char.from_buffer(mv))
def to_char_p_p(options: list[bytes], to_type=ctypes.c_char):
return (ctypes.POINTER(to_type) * len(options))(*[ctypes.cast(ctypes.create_string_buffer(o), ctypes.POINTER(to_type)) for o in options])
def charptr(s:str|bytes): return ctypes.cast(ctypes.c_char_p(s if isinstance(s, bytes) else s.encode()), ctypes.POINTER(ctypes.c_char))
@functools.cache
def init_c_struct_t(fields: tuple[tuple[str, type[ctypes._SimpleCData]], ...]):
class CStruct(ctypes.Structure):
tinygrad/mixin/__init__.py
+3 -1
View File
@@ -1,4 +1,6 @@
from tinygrad.mixin.math import MathMixin
from tinygrad.mixin.movement import MovementMixin
class OpMixin(MathMixin, MovementMixin): pass
class OpMixin(MathMixin, MovementMixin):
pass
tinygrad/mixin/math.py
+173 -66
View File
@@ -2,24 +2,38 @@ from typing import Self
from tinygrad.uop import Ops
from tinygrad.dtype import dtypes, ConstType
class MathMixin:
# required to implement
def alu(self, op:Ops, *src:Self) -> Self: raise NotImplementedError
def const_like(self, b:ConstType) -> Self: raise NotImplementedError
def alu(self, op: Ops, *src: Self) -> Self:
raise NotImplementedError
def const_like(self, b: ConstType) -> Self:
raise NotImplementedError
# great functions you get!
def ufix(self, x:Self|ConstType) -> Self: return self.const_like(x) if not isinstance(x, MathMixin) else x
def _binop(self, op:Ops, x:Self|ConstType, reverse:bool) -> Self:
def ufix(self, x: Self | ConstType) -> Self:
return self.const_like(x) if not isinstance(x, MathMixin) else x
def _binop(self, op: Ops, x: Self | ConstType, reverse: bool) -> Self:
return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x))
def logical_not(self): return self.ne(True)
def logical_not(self):
return self.ne(True)
def neg(self):
if (dtype:=getattr(self, 'dtype')) is None: raise TypeError(f"MathTraits __neg__ requires a dtype, {self=}")
return self.logical_not() if dtype.scalar() == dtypes.bool else self*(-1)
if (dtype := getattr(self, "dtype")) is None:
raise TypeError(f"MathTraits __neg__ requires a dtype, {self=}")
return self.logical_not() if dtype.scalar() == dtypes.bool else self * (-1)
def _check_dtype(self):
if (dtype:=getattr(self, 'dtype')) is not None:
if isinstance(dtype, tuple): dtype = dtype[0]
if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)): raise RuntimeError(f"{dtype} is not supported")
def add(self, x:Self|ConstType, reverse:bool=False):
if (dtype := getattr(self, "dtype")) is not None:
if isinstance(dtype, tuple):
dtype = dtype[0]
if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)):
raise RuntimeError(f"{dtype} is not supported")
def add(self, x: Self | ConstType, reverse: bool = False):
"""
Adds `self` and `x`.
Equivalent to `self + x`.
@@ -37,7 +51,8 @@ class MathMixin:
```
"""
return self._binop(Ops.ADD, x, reverse)
def mul(self, x:Self|ConstType, reverse:bool=False):
def mul(self, x: Self | ConstType, reverse: bool = False):
"""
Multiplies `self` and `x`.
Equivalent to `self * x`.
@@ -56,7 +71,8 @@ class MathMixin:
```
"""
return self._binop(Ops.MUL, x, reverse)
def bitwise_and(self, x:Self|ConstType, reverse:bool=False):
def bitwise_and(self, x: Self | ConstType, reverse: bool = False):
"""
Computes the bitwise AND of `self` and `x`.
Equivalent to `self & x`.
@@ -70,7 +86,8 @@ class MathMixin:
"""
self._check_dtype()
return self._binop(Ops.AND, x, reverse)
def bitwise_or(self, x:Self|ConstType, reverse:bool=False):
def bitwise_or(self, x: Self | ConstType, reverse: bool = False):
"""
Computes the bitwise OR of `self` and `x`.
Equivalent to `self | x`.
@@ -84,7 +101,8 @@ class MathMixin:
"""
self._check_dtype()
return self._binop(Ops.OR, x, reverse)
def bitwise_xor(self, x:Self|ConstType, reverse:bool=False):
def bitwise_xor(self, x: Self | ConstType, reverse: bool = False):
"""
Computes bitwise xor of `self` and `x`.
Equivalent to `self ^ x`.
@@ -99,7 +117,8 @@ class MathMixin:
"""
self._check_dtype()
return self._binop(Ops.XOR, x, reverse)
def idiv(self, x:Self|ConstType, reverse:bool=False):
def idiv(self, x: Self | ConstType, reverse: bool = False):
"""
Divides `self` by `x`.
Equivalent to `self // x`.
@@ -111,62 +130,150 @@ class MathMixin:
```
"""
return self._binop(Ops.IDIV, x, reverse)
def mod(self, x:Self|ConstType, reverse:bool=False): return self._binop(Ops.MOD, x, reverse)
def sub(self, x:Self|ConstType, reverse:bool=False): return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x))
def div(self, x:Self|ConstType, reverse:bool=False):
return (self.ufix(x)*self.alu(Ops.RECIPROCAL)) if reverse else (self*self.ufix(x).alu(Ops.RECIPROCAL))
def __neg__(self): return self.neg()
def mod(self, x: Self | ConstType, reverse: bool = False):
return self._binop(Ops.MOD, x, reverse)
def __add__(self, x:Self|ConstType): return self.add(x)
def __sub__(self, x:Self|ConstType): return self.sub(x)
def __mul__(self, x:Self|ConstType): return self.mul(x)
def __truediv__(self, x:Self|ConstType): return self.div(x)
def __floordiv__(self, x:Self|ConstType): return self.idiv(x) # TODO: idiv is trunc div, not floordiv
def __mod__(self, x:Self|ConstType): return self.mod(x)
def __and__(self, x:Self|ConstType): return self.bitwise_and(x)
def __or__(self, x:Self|ConstType): return self.bitwise_or(x)
def __xor__(self, x:Self|ConstType): return self.bitwise_xor(x)
def sub(self, x: Self | ConstType, reverse: bool = False):
return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x))
def __radd__(self, x:Self|ConstType): return self.add(x, True)
def __rsub__(self, x:Self|ConstType): return self.sub(x, True)
def __rmul__(self, x:Self|ConstType): return self.mul(x, True)
def __rtruediv__(self, x:Self|ConstType): return self.div(x, True)
def __rfloordiv__(self, x:Self|ConstType): return self.idiv(x, True)
def __rand__(self, x:Self|ConstType): return self.bitwise_and(x, True)
def __ror__(self, x:Self|ConstType): return self.bitwise_or(x, True)
def __rxor__(self, x:Self|ConstType): return self.bitwise_xor(x, True)
def __rmod__(self, x:Self|ConstType): return self.mod(x, True)
def div(self, x: Self | ConstType, reverse: bool = False):
return (self.ufix(x) * self.alu(Ops.RECIPROCAL)) if reverse else (self * self.ufix(x).alu(Ops.RECIPROCAL))
def __lt__(self, x:Self|ConstType): return self.alu(Ops.CMPLT, self.ufix(x))
def __gt__(self, x:Self|ConstType): return self.ufix(x).alu(Ops.CMPLT, self)
def __ge__(self, x:Self|ConstType): return (self < x).logical_not()
def __le__(self, x:Self|ConstType): return (self > x).logical_not()
def __neg__(self):
return self.neg()
def __add__(self, x: Self | ConstType):
return self.add(x)
def __sub__(self, x: Self | ConstType):
return self.sub(x)
def __mul__(self, x: Self | ConstType):
return self.mul(x)
def __truediv__(self, x: Self | ConstType):
return self.div(x)
def __floordiv__(self, x: Self | ConstType):
return self.idiv(x) # TODO: idiv is trunc div, not floordiv
def __mod__(self, x: Self | ConstType):
return self.mod(x)
def __and__(self, x: Self | ConstType):
return self.bitwise_and(x)
def __or__(self, x: Self | ConstType):
return self.bitwise_or(x)
def __xor__(self, x: Self | ConstType):
return self.bitwise_xor(x)
def __radd__(self, x: Self | ConstType):
return self.add(x, True)
def __rsub__(self, x: Self | ConstType):
return self.sub(x, True)
def __rmul__(self, x: Self | ConstType):
return self.mul(x, True)
def __rtruediv__(self, x: Self | ConstType):
return self.div(x, True)
def __rfloordiv__(self, x: Self | ConstType):
return self.idiv(x, True)
def __rand__(self, x: Self | ConstType):
return self.bitwise_and(x, True)
def __ror__(self, x: Self | ConstType):
return self.bitwise_or(x, True)
def __rxor__(self, x: Self | ConstType):
return self.bitwise_xor(x, True)
def __rmod__(self, x: Self | ConstType):
return self.mod(x, True)
def __lt__(self, x: Self | ConstType):
return self.alu(Ops.CMPLT, self.ufix(x))
def __gt__(self, x: Self | ConstType):
return self.ufix(x).alu(Ops.CMPLT, self)
def __ge__(self, x: Self | ConstType):
return (self < x).logical_not()
def __le__(self, x: Self | ConstType):
return (self > x).logical_not()
def ne(self, x: Self | ConstType):
return self.alu(Ops.CMPNE, self.ufix(x))
def eq(self, x: Self | ConstType):
return self.ne(x).logical_not()
def __ne__(self, x: Self | ConstType): # type: ignore[override]
return self.ne(x)
def ne(self, x:Self|ConstType): return self.alu(Ops.CMPNE, self.ufix(x))
def eq(self, x:Self|ConstType): return self.ne(x).logical_not()
def __ne__(self, x:Self|ConstType): return self.ne(x) # type: ignore[override]
# NOTE: __eq__ isn't overridden, and means the same thing as is by default
def lshift(self, x:Self|int, reverse:bool=False): return self._binop(Ops.SHL, x, reverse)
def rshift(self, x:Self|int, reverse:bool=False): return self._binop(Ops.SHR, x, reverse)
def __lshift__(self, x:Self|int): return self.lshift(x)
def __rshift__(self, x:Self|int): return self.rshift(x)
def __rlshift__(self, x:Self|int): return self.lshift(x, True)
def __rrshift__(self, x:Self|int): return self.rshift(x, True)
def lshift(self, x: Self | int, reverse: bool = False):
return self._binop(Ops.SHL, x, reverse)
def maximum(self, x:Self|ConstType): return self.alu(Ops.MAX, self.ufix(x))
def minimum(self, x:Self|ConstType): return -(-self).maximum(-x)
def where(self, x:Self|ConstType, y:Self|ConstType):
if isinstance(x, type(self)): return self.alu(Ops.WHERE, x, x.ufix(y))
if isinstance(y, type(self)): return self.alu(Ops.WHERE, y.ufix(x), y)
def rshift(self, x: Self | int, reverse: bool = False):
return self._binop(Ops.SHR, x, reverse)
def __lshift__(self, x: Self | int):
return self.lshift(x)
def __rshift__(self, x: Self | int):
return self.rshift(x)
def __rlshift__(self, x: Self | int):
return self.lshift(x, True)
def __rrshift__(self, x: Self | int):
return self.rshift(x, True)
def maximum(self, x: Self | ConstType):
return self.alu(Ops.MAX, self.ufix(x))
def minimum(self, x: Self | ConstType):
return -(-self).maximum(-x)
def where(self, x: Self | ConstType, y: Self | ConstType):
if isinstance(x, type(self)):
return self.alu(Ops.WHERE, x, x.ufix(y))
if isinstance(y, type(self)):
return self.alu(Ops.WHERE, y.ufix(x), y)
raise RuntimeError("where needs at least one UOp arg")
def threefry(self, seed:Self): return self.alu(Ops.THREEFRY, seed)
def reciprocal(self): return self.alu(Ops.RECIPROCAL)
def trunc(self): return self.alu(Ops.TRUNC)
def sqrt(self): return self.alu(Ops.SQRT)
def sin(self): return self.alu(Ops.SIN)
def log2(self): return self.alu(Ops.LOG2)
def exp2(self): return self.alu(Ops.EXP2)
def pow(self, x:Self|ConstType): return self.alu(Ops.POW, self.ufix(x))
def __pow__(self, x:Self|ConstType): return self.pow(x)
def threefry(self, seed: Self):
return self.alu(Ops.THREEFRY, seed)
def reciprocal(self):
return self.alu(Ops.RECIPROCAL)
def trunc(self):
return self.alu(Ops.TRUNC)
def sqrt(self):
return self.alu(Ops.SQRT)
def sin(self):
return self.alu(Ops.SIN)
def log2(self):
return self.alu(Ops.LOG2)
def exp2(self):
return self.alu(Ops.EXP2)
def pow(self, x: Self | ConstType):
return self.alu(Ops.POW, self.ufix(x))
def __pow__(self, x: Self | ConstType):
return self.pow(x)
tinygrad/mixin/movement.py
+88 -40
View File
@@ -2,20 +2,28 @@
import functools
from typing import TypeAlias, TYPE_CHECKING, Self
from tinygrad.uop import Ops
from tinygrad.helpers import prod, argfix, flatten, dedup
if TYPE_CHECKING: from tinygrad.uop.ops import UOp
from tinygrad.helpers import prod, argfix, flatten, dedup, make_tuple, ceildiv
from tinygrad.uop.ops import resolve, smax
if TYPE_CHECKING:
from tinygrad.uop.ops import UOp
sint: TypeAlias = "UOp | int"
def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
def _align_left(*shapes: tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
# unsqueeze left to make every shape same length
max_dim = max(len(shape) for shape in shapes)
return tuple((1,) * (max_dim - len(shape)) + shape for shape in shapes)
class MovementMixin:
# required to implement
def _mop(self, op:Ops, arg) -> Self: raise NotImplementedError
def _mop(self, op: Ops, arg) -> Self:
raise NotImplementedError
@property
def shape(self) -> tuple[sint, ...]: raise NotImplementedError
def shape(self) -> tuple[sint, ...]:
raise NotImplementedError
# great functions you get!
@property
@@ -41,18 +49,21 @@ class MovementMixin:
"""
return prod(self.shape)
def _resolve_dim(self, dim:int, *, extra:bool=False) -> int:
def _resolve_dim(self, dim: int, *, extra: bool = False) -> int:
total = self.ndim + int(extra)
if not -max(1, total) <= dim <= max(1, total)-1: raise IndexError(f"{dim=} out of range {[-max(1, total), max(1, total)-1]}")
if not -max(1, total) <= dim <= max(1, total) - 1:
raise IndexError(f"{dim=} out of range {[-max(1, total), max(1, total) - 1]}")
return dim + total if dim < 0 else dim
def _broadcast_to(self, new_shape:tuple[sint, ...]) -> Self:
if self.shape == new_shape: return self
if self.ndim > len(new_shape): raise ValueError(f"cannot broadcast tensor to fewer dimensions. shape={self.shape} to {new_shape=}")
def _broadcast_to(self, new_shape: tuple[sint, ...]) -> Self:
if self.shape == new_shape:
return self
if self.ndim > len(new_shape):
raise ValueError(f"cannot broadcast tensor to fewer dimensions. shape={self.shape} to {new_shape=}")
# first unsqueeze left with 1s https://data-apis.org/array-api/latest/API_specification/broadcasting.html
shape, _ = _align_left(self.shape, new_shape)
# for each dimension, check either dim is 1, or it does not change
if not all(s == ns or s == 1 for s,ns in zip(shape, new_shape)):
if not all(s == ns or s == 1 for s, ns in zip(shape, new_shape)):
raise ValueError(f"cannot broadcast {self.shape} to {new_shape=}")
reshaped = self.reshape(shape)
ret = reshaped._mop(Ops.EXPAND, arg=new_shape)
@@ -84,15 +95,18 @@ class MovementMixin:
```
"""
# resolve None and args
new_shape = tuple([s if s is not None else self.shape[i] for i,s in enumerate(argfix(shape, *args))])
new_shape = tuple([s if s is not None else self.shape[i] for i, s in enumerate(argfix(shape, *args))])
# resolve -1
if (c := new_shape.count(-1)) > 1: raise RuntimeError(f"only one dimension can be inferred using -1, getting {new_shape}")
if c: new_shape = tuple([-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape])
if prod(self.shape) != prod(new_shape): raise ValueError(f"size mismatch, can't reshape ({self.shape}) -> ({new_shape})")
if (c := new_shape.count(-1)) > 1:
raise RuntimeError(f"only one dimension can be inferred using -1, getting {new_shape}")
if c:
new_shape = tuple([-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape])
if prod(self.shape) != prod(new_shape):
raise ValueError(f"size mismatch, can't reshape ({self.shape}) -> ({new_shape})")
ret = self._mop(Ops.RESHAPE, arg=new_shape)
return self if ret.shape == self.shape else ret
def shrink(self, arg:tuple[tuple[sint, sint]|None, ...]) -> Self:
def shrink(self, arg: tuple[tuple[sint, sint] | None, ...]) -> Self:
"""
Returns a tensor that shrinks the each axis based on input arg.
`arg` must have the same length as `self.ndim`.
@@ -109,8 +123,9 @@ class MovementMixin:
print(t.shrink((((0, 2), (0, 2)))).numpy())
```
"""
if self.ndim != len(arg): raise ValueError(f"{self.ndim=} != {len(arg)=}")
ret = self._mop(Ops.SHRINK, arg=[x if x is not None else (0,s) for x,s in zip(arg, self.shape)])
if self.ndim != len(arg):
raise ValueError(f"{self.ndim=} != {len(arg)=}")
ret = self._mop(Ops.SHRINK, arg=[x if x is not None else (0, s) for x, s in zip(arg, self.shape)])
return self if ret.shape == self.shape else ret
def permute(self, order, *args) -> Self:
@@ -128,7 +143,8 @@ class MovementMixin:
```
"""
order_arg = tuple(self._resolve_dim(x) for x in argfix(order, *args))
if sorted(order_arg) != list(range(self.ndim)): raise RuntimeError(f"order is not a valid permutation, getting {order_arg}")
if sorted(order_arg) != list(range(self.ndim)):
raise RuntimeError(f"order is not a valid permutation, getting {order_arg}")
return self._mop(Ops.PERMUTE, arg=order_arg) if order_arg != tuple(range(self.ndim)) else self
def flip(self, axis, *args) -> Self:
@@ -149,7 +165,8 @@ class MovementMixin:
"""
axis_arg = tuple(self._resolve_dim(x) for x in argfix(axis, *args))
assert all(not isinstance(x, bool) and x >= 0 and x < self.ndim for x in axis_arg), f"flip args must be axis ints {axis_arg}"
if len(axis_arg) != len(dedup(axis_arg)): raise RuntimeError(f"dim can appear at most once, getting {axis_arg}")
if len(axis_arg) != len(dedup(axis_arg)):
raise RuntimeError(f"dim can appear at most once, getting {axis_arg}")
flip_arg = tuple([i in axis_arg for i in range(len(self.shape))])
return self._mop(Ops.FLIP, arg=flip_arg) if any(flip_arg) else self
@@ -162,7 +179,7 @@ class MovementMixin:
"""`.view` is an alias for `.reshape`."""
return self.reshape(shape, *args)
def squeeze(self, dim:int|None=None) -> Self:
def squeeze(self, dim: int | None = None) -> Self:
"""
Returns a tensor with specified dimensions of input of size 1 removed.
If `dim` is not specified, all dimensions with size 1 are removed.
@@ -178,11 +195,12 @@ class MovementMixin:
print(t.squeeze(1).shape)
```
"""
if dim is None: return self.reshape(tuple(dim for dim in self.shape if dim != 1))
if dim is None:
return self.reshape(tuple(dim for dim in self.shape if dim != 1))
dim = self._resolve_dim(dim)
return self if not self.ndim or self.shape[dim] != 1 else self.reshape(self.shape[:dim] + self.shape[dim+1:])
return self if not self.ndim or self.shape[dim] != 1 else self.reshape(self.shape[:dim] + self.shape[dim + 1 :])
def unsqueeze(self, dim:int) -> Self:
def unsqueeze(self, dim: int) -> Self:
"""
Returns a tensor with a new dimension of size 1 inserted at the specified `dim`.
@@ -233,9 +251,9 @@ class MovementMixin:
```
"""
start_dim, end_dim = self._resolve_dim(start_dim), self._resolve_dim(end_dim)
return self.reshape(self.shape[:start_dim] + (prod(self.shape[start_dim:end_dim+1]), ) + self.shape[end_dim+1:])
return self.reshape(self.shape[:start_dim] + (prod(self.shape[start_dim : end_dim + 1]),) + self.shape[end_dim + 1 :])
def unflatten(self, dim:int, sizes:tuple[int,...]) -> Self:
def unflatten(self, dim: int, sizes: tuple[int, ...]) -> Self:
"""
Unflattens dimension `dim` of the tensor into multiple dimensions specified by `sizes`. `Tensor.flatten()` is the inverse of this function.
@@ -250,9 +268,9 @@ class MovementMixin:
```
"""
dim = self._resolve_dim(dim)
return self.reshape(self.shape[:dim] + sizes + self.shape[dim+1:])
return self.reshape(self.shape[:dim] + sizes + self.shape[dim + 1 :])
def rearrange(self, formula:str, **sizes) -> Self:
def rearrange(self, formula: str, **sizes) -> Self:
"""
Rearranges input according to formula
@@ -263,38 +281,43 @@ class MovementMixin:
print(Tensor.rearrange(x, "batch channel -> (batch channel)").numpy())
```
"""
def parse_formula(formula: str):
tokens = f" {formula} ".replace("…", "...").replace("(", " ( ").replace(")", " ) ").replace(" ", " ").replace(" 1 ", " ( ) ").split()
lparens, rparens = map(lambda x: [i for i, ch in enumerate(tokens) if ch == x], ("(", ")"))
pairs = list(zip(lparens, rparens))
assert len(lparens) == len(rparens) and sorted(flatten(pairs)) == flatten(pairs), "bracket mismatch"
return [name for name in tokens if name not in ("(", ")")], [(s - 2*i, e - 1 - 2*i) for i, (s, e) in enumerate(pairs)]
return [name for name in tokens if name not in ("(", ")")], [(s - 2 * i, e - 1 - 2 * i) for i, (s, e) in enumerate(pairs)]
assert formula.count("->") == 1, 'need exactly one "->" in formula'
(lhs, unflatten_dims), (rhs, flatten_dims) = map(parse_formula, formula.split("->"))
for name in sizes: assert name in lhs, f"axis {name} is not used in transform"
for name in sizes:
assert name in lhs, f"axis {name} is not used in transform"
assert sorted(lhs) == sorted(rhs) and len(lhs) == len(set(lhs)), f"name mismatch in {formula}"
for name in flatten((lhs, rhs)): assert name == "..." or (name.isidentifier() and "_" not in (name[0], name[-1])), f"invalid axis name {name}"
for name in flatten((lhs, rhs)):
assert name == "..." or (name.isidentifier() and "_" not in (name[0], name[-1])), f"invalid axis name {name}"
assert "..." not in flatten([lhs[s:e] for s, e in unflatten_dims]), f"cannot have collapsed ellipsis (...) in lhs of {formula}"
assert lhs.count("...") <= 1, f"too many ellipses in {formula}"
# resolve ellipsis
if "..." in lhs: ell_len = len(self.shape) - len(lhs) + 1 + sum(e - s - 1 for s, e in unflatten_dims)
lhs, rhs = map(lambda l: l[:(i:=l.index("..."))] + [f"...{j}" for j in range(ell_len)] + l[i + 1:] if "..." in l else l, (lhs, rhs))
if "..." in lhs:
ell_len = len(self.shape) - len(lhs) + 1 + sum(e - s - 1 for s, e in unflatten_dims)
lhs, rhs = map(lambda l: l[: (i := l.index("..."))] + [f"...{j}" for j in range(ell_len)] + l[i + 1 :] if "..." in l else l, (lhs, rhs))
unflatten_dims = [(s + (ell_len - 1 if "...0" in lhs[:s] else 0), e + (ell_len - 1 if "...0" in lhs[:e] else 0)) for s, e in unflatten_dims]
flatten_dims = [(s + (ell_len - 1 if "...0" in rhs[:s] else 0), e + (ell_len - 1 if "...0" in rhs[:e] else 0)) for s, e in flatten_dims]
# apply movement ops in order unflatten -> permute -> flatten/unsqueeze
t = functools.reduce(lambda x, dims: x.unflatten(dims[0], tuple(sizes.get(lhs[d], -1) for d in range(*dims))), unflatten_dims, self)
for i, name in enumerate(lhs): assert (name not in sizes) or sizes[name] == t.shape[i], f"size provided for dimension {name} incorrect"
for i, name in enumerate(lhs):
assert (name not in sizes) or sizes[name] == t.shape[i], f"size provided for dimension {name} incorrect"
t = t.permute([lhs.index(name) for name in rhs])
return functools.reduce(lambda x, dims: x.flatten(dims[0], dims[1] - 1) if dims[0]<dims[1] else x.unsqueeze(dims[0]), reversed(flatten_dims), t)
return functools.reduce(lambda x, dims: x.flatten(dims[0], dims[1] - 1) if dims[0] < dims[1] else x.unsqueeze(dims[0]), reversed(flatten_dims), t)
# *** movement ops with expand ***
def repeat_interleave(self, repeats:int, dim:int|None=None) -> Self:
def repeat_interleave(self, repeats: int, dim: int | None = None) -> Self:
"""
Repeats elements of a tensor.
@@ -305,7 +328,10 @@ class MovementMixin:
"""
x, dim = (self.flatten(), 0) if dim is None else (self, self._resolve_dim(dim))
shp = x.shape
return x.reshape(*shp[:dim+1], 1, *shp[dim+1:]).expand(*shp[:dim+1], repeats, *shp[dim+1:]).reshape(*shp[:dim], shp[dim]*repeats, *shp[dim+1:])
x = x.reshape(*shp[: dim + 1], 1, *shp[dim + 1 :])
x = x.expand(*shp[: dim + 1], repeats, *shp[dim + 1 :])
x = x.reshape(*shp[:dim], shp[dim] * repeats, *shp[dim + 1 :])
return x
def repeat(self, repeats, *args) -> Self:
"""
@@ -322,7 +348,29 @@ class MovementMixin:
"""
repeats = argfix(repeats, *args)
base_shape = _align_left(self.shape, repeats)[0]
unsqueezed_shape = flatten([[s] if r == 1 else [1, s] for r,s in zip(repeats, base_shape)])
expanded_shape = flatten([[s] if r == 1 else [r, s] for r,s in zip(repeats, base_shape)])
final_shape = [r*s for r,s in zip(repeats, base_shape)]
unsqueezed_shape = flatten([[s] if r == 1 else [1, s] for r, s in zip(repeats, base_shape)])
expanded_shape = flatten([[s] if r == 1 else [r, s] for r, s in zip(repeats, base_shape)])
final_shape = [r * s for r, s in zip(repeats, base_shape)]
return self.reshape(unsqueezed_shape).expand(expanded_shape).reshape(final_shape)
# **** pool level ****
def _pool(self, k_: tuple[sint, ...], stride: int | tuple[int, ...] = 1, dilation: int | tuple[int, ...] = 1) -> Self:
assert len(self.shape) >= len(k_), f"can't pool {self.shape} with {k_}"
s_, d_ = make_tuple(stride, len(k_)), make_tuple(dilation, len(k_))
assert len(k_) == len(s_) == len(d_), f"stride/dilation mismatch kernel:{k_} stride:{s_} dilation:{d_}"
noop, i_ = [None] * (self.ndim - len(k_)), self.shape[-len(k_) :]
assert all(resolve(d * (k - 1) + 1 <= i) for k, d, i in zip(k_, d_, i_)), "kernel size cannot be greater than actual input size"
o_ = [ceildiv(i - d * (k - 1), s) for i, d, k, s in zip(i_, d_, k_, s_)]
# input size scaling factor to make sure shrink for stride is possible
f_ = [smax(1, ceildiv(o * s - d, i)) for o, s, i, d in zip(o_, s_, i_, d_)]
# repeats such that we don't need padding
x = self.repeat([1] * len(noop) + [ceildiv(k * (i * f + d), i) for k, i, d, f in zip(k_, i_, d_, f_)])
# handle dilation
x = x.shrink_to(noop + [k * (i * f + d) for k, i, d, f in zip(k_, i_, d_, f_)])
x = x.reshape(noop + flatten((k, (i * f + d)) for k, i, d, f in zip(k_, i_, d_, f_)))
# handle stride
x = x.shrink_to(noop + flatten((k, o * s) for k, o, s in zip(k_, o_, s_))).reshape(noop + flatten((k, o, s) for k, o, s in zip(k_, o_, s_)))
x = x.shrink_to(noop + flatten((k, o, 1) for k, o in zip(k_, o_))).reshape(noop + flatten((k, o) for k, o in zip(k_, o_)))
# permute to move reduce to the end
return x.permute(*range(len(noop)), *[len(noop) + i * 2 + 1 for i in range(len(i_))], *[len(noop) + i * 2 for i in range(len(i_))])
tinygrad/nn/onnx.py
+10
View File
@@ -1124,6 +1124,16 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
return output.flatten(start_dim=2) if len(original_input_shape) == 3 else output.permute(0, 2, 1, 3)
# ***** Indexing Ops *****
def NonZero(x:Tensor):
mask = (x!=0).flatten()
flat_idx = Tensor.arange(mask.numel(), dtype=dtypes.int64, device=x.device).masked_select(mask)
if flat_idx.ndim == 0: flat_idx = flat_idx.reshape(1)
if x.ndim == 0:
return Tensor.zeros((0, flat_idx.shape[0]), dtype=dtypes.int64, device=x.device, requires_grad=False)
strides = [prod(int(s) for s in x.shape[i+1:]) if i+1 < x.ndim else 1 for i in range(x.ndim)]
coords = [((flat_idx // stride) % int(dim)) for stride, dim in zip(strides, x.shape)]
return Tensor.stack(*coords, dim=0)
def ArrayFeatureExtractor(x:Tensor, indices:Tensor): return x[..., indices]
def Gather(x:Tensor, indices:Tensor, axis:int=0):
tinygrad/nn/state.py
+6 -1
View File
@@ -194,6 +194,10 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:
"""
offsets: dict[str|int, int] = {}
lens: dict[str|int, int] = {}
def _rebuild_tensor(storage, storage_offset, size, stride):
return _rebuild_tensor_v2(storage, storage_offset, size, stride)
def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad=None, backward_hooks=None, metadata=None):
#print(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata)
lens[storage[2]] = storage[4] * storage[1].itemsize
@@ -220,7 +224,8 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:
deserialized_objects: dict[str, Any] = {}
intercept = {"HalfStorage": dtypes.float16, "FloatStorage": dtypes.float32, "BFloat16Storage": dtypes.bfloat16,
"IntStorage": dtypes.int32, "BoolStorage": dtypes.bool,
"LongStorage": dtypes.int64, "_rebuild_tensor_v2": _rebuild_tensor_v2, "FloatTensor": None, "Parameter": Parameter}
"LongStorage": dtypes.int64, "_rebuild_tensor": _rebuild_tensor, "_rebuild_tensor_v2": _rebuild_tensor_v2,
"FloatTensor": None, "Parameter": Parameter}
whitelist = {"torch", "collections", "numpy", "_codecs"} # NOTE: this is not for security, only speed
class Dummy: pass
class TorchPickle(pickle.Unpickler):
tinygrad/renderer/nir.py
+4 -4
View File
@@ -1,10 +1,10 @@
from typing import Callable, cast, Any
from tinygrad.dtype import AddrSpace, DType, PtrDType, dtypes
from tinygrad.helpers import DEBUG, OSX, unwrap
from tinygrad.helpers import DEBUG, OSX, unwrap, charptr
from tinygrad.renderer import Renderer
from tinygrad.renderer.cstyle import CUDARenderer
from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
import tinygrad.runtime.autogen.mesa as mesa
from tinygrad.runtime.autogen import mesa
import base64, ctypes, ctypes.util, struct, functools, inspect
def g(s:str): return getattr(mesa, s)
@@ -51,7 +51,7 @@ def nir_instr(nc=1, bs=lambda: None, intrins=None, srcs=None, has_def=True, df=N
instr = f(*args, **kwargs)
if has_def: mesa.nir_def_init(instr.contents.instr, getattr(instr.contents, "def"), go(nc), go(bs))
for k, v in go(intrins or {}).items():
idx = mesa.nir_intrinsic_infos[instr.contents.intrinsic].index_map[g(f"NIR_INTRINSIC_{k}")]
idx = mesa.nir_intrinsic_infos[instr.contents.intrinsic.value].index_map[g(f"NIR_INTRINSIC_{k}")]
assert idx > 0
instr.contents.const_index[idx - 1] = go(v)
for i, src in enumerate(go(srcs or [])): ctypes.cast(instr.contents.src, ctypes.POINTER(mesa.nir_src))[i] = go(src)
@@ -177,7 +177,7 @@ class NIRRenderer(Renderer):
elif u.op is Ops.AFTER:
self.r[u] = self.r[u.src[0]]
elif u.op == Ops.SINK:
if u.arg is not None: self.b.shader.contents.info.name = mesa.char_pointer_cast(u.arg.function_name)
if u.arg is not None: self.b.shader.contents.info.name = charptr(u.arg.function_name.encode())
elif u.op == Ops.DEFINE_LOCAL:
self.r[u] = nimm(self.b, self.b.shader.contents.info.shared_size, dtypes.long)
self.b.shader.contents.info.shared_size += u.dtype.nbytes()
tinygrad/runtime/autogen/__init__.py
+109 -3
View File
@@ -1,11 +1,23 @@
import importlib, pathlib
from tinygrad.helpers import system
import glob, importlib, pathlib, subprocess, tarfile
from tinygrad.helpers import fetch, flatten, system
root = (here:=pathlib.Path(__file__).parent).parents[2]
nv_src = {"nv_570": "https://github.com/NVIDIA/open-gpu-kernel-modules/archive/81fe4fb417c8ac3b9bdcc1d56827d116743892a5.tar.gz",
"nv_580": "https://github.com/NVIDIA/open-gpu-kernel-modules/archive/2af9f1f0f7de4988432d4ae875b5858ffdb09cc2.tar.gz"}
def load(name, dll, files, **kwargs):
if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists():
files = files() if callable(files) else files
files, kwargs['args'] = files() if callable(files) else files, args() if callable(args:=kwargs.get('args', [])) else args
if (tarball:=kwargs.pop('tarball', None)):
# dangerous for arbitrary urls!
with tarfile.open(fetch(tarball, gunzip=tarball.endswith("gz"))) as tf:
tf.extractall("/tmp")
base = f"/tmp/{tf.getnames()[0]}"
files, kwargs['args'] = [str(f).format(base) for f in files], [a.format(base) for a in kwargs.get('args', [])]
kwargs['anon_names'] = {k.format(base):v for k,v in kwargs.get('anon_names', {}).items()}
if (preprocess:=kwargs.pop('preprocess', None)): preprocess(base)
files = flatten(sorted(glob.glob(p, recursive=True)) if isinstance(p, str) and '*' in p else [p] for p in files)
kwargs['epilog'] = (epi(base) if tarball else epi()) if callable(epi:=kwargs.get('epilog', [])) else epi
f.write_text(importlib.import_module("tinygrad.runtime.support.autogen").gen(dll, files, **kwargs))
return importlib.import_module(f"{path}.{name.replace('/', '.')}")
@@ -14,4 +26,98 @@ def __getattr__(nm):
case "libc": return load("libc", ["find_library('c')"], lambda: (
[i for i in system("dpkg -L libc6-dev").split() if 'sys/mman.h' in i or 'sys/syscall.h' in i] +
["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/asm-generic/mman-common.h"]), use_errno=True)
case "opencl": return load("opencl", ["find_library('OpenCL')"], ["/usr/include/CL/cl.h"])
case "cuda": return load("cuda", ["find_library('cuda')"], ["/usr/include/cuda.h"], args=["-D__CUDA_API_VERSION_INTERNAL"], parse_macros=False)
case "nvrtc": return load("nvrtc", ["find_library('nvrtc')"], ["/usr/include/nvrtc.h"])
case "nvjitlink": load("nvjitlink", ["find_library('nvJitLink')"], [root/"extra/nvJitLink.h"])
case "kfd": return load("kfd", [], ["/usr/include/linux/kfd_ioctl.h"])
case "nv_570" | "nv_580":
return load(nm, [], [
*[root/"extra/nv_gpu_driver"/s for s in ["clc6c0qmd.h","clcec0qmd.h"]], "{}/kernel-open/common/inc/nvmisc.h",
*[f"{{}}/src/common/sdk/nvidia/inc/class/cl{s}.h" for s in ["0000", "0080", "2080", "2080_notification", "c56f", "c86f", "c96f", "c761",
"83de", "c6c0", "cdc0"]],
*[f"{{}}/kernel-open/nvidia-uvm/{s}.h" for s in ["clc6b5", "clc9b5", "uvm_ioctl", "uvm_linux_ioctl", "hwref/ampere/ga100/dev_fault"]],
*[f"{{}}/src/nvidia/arch/nvalloc/unix/include/nv{s}.h" for s in ["_escape", "-ioctl", "-ioctl-numbers",
"-ioctl-numa", "-unix-nvos-params-wrappers"]],
*[f"{{}}/src/common/sdk/nvidia/inc/{s}.h" for s in ["alloc/alloc_channel", "nvos", "ctrl/ctrlc36f", "ctrl/ctrlcb33",
"ctrl/ctrla06c", "ctrl/ctrl90f1"]],
*[f"{{}}/src/common/sdk/nvidia/inc/ctrl/ctrl{s}/*.h" for s in ["0000", "0080", "2080", "83de"]],
"{}/kernel-open/common/inc/nvstatus.h", "{}/src/nvidia/generated/g_allclasses.h"
], args=[
"-include", "{}/src/common/sdk/nvidia/inc/nvtypes.h", "-I{}/src/common/inc", "-I{}/kernel-open/nvidia-uvm", "-I{}/kernel-open/common/inc",
"-I{}/src/common/sdk/nvidia/inc", "-I{}/src/nvidia/arch/nvalloc/unix/include", "-I{}/src/common/sdk/nvidia/inc/ctrl"
], rules=[(r'MW\(([^:]+):(.+)\)',r'(\1, \2)')], tarball=nv_src[nm], anon_names={"{}/kernel-open/common/inc/nvstatus.h:37":"nv_status_codes"})
case "nv": return load("nv", [], [
*[f"{{}}/src/nvidia/inc/kernel/gpu/{s}.h" for s in ["fsp/kern_fsp_cot_payload", "gsp/gsp_init_args"]],
*[f"{{}}/src/nvidia/arch/nvalloc/common/inc/{s}.h" for s in ["gsp/gspifpub", "gsp/gsp_fw_wpr_meta", "gsp/gsp_fw_sr_meta", "rmRiscvUcode",
"fsp/fsp_nvdm_format"]],
*[f"{{}}/src/nvidia/inc/kernel/vgpu/{s}.h" for s in ["rpc_headers", "rpc_global_enums"]],
"{}/src/common/uproc/os/common/include/libos_init_args.h", "{}/src/common/shared/msgq/inc/msgq/msgq_priv.h",
"{}/src/nvidia/generated/g_rpc-structures.h", root/"extra/nv_gpu_driver/g_rpc-message-header.h", root/"extra/nv_gpu_driver/gsp_static_config.h",
root/"extra/nv_gpu_driver/vbios.h", root/"extra/nv_gpu_driver/pci_exp_table.h"
], args=[
"-DRPC_MESSAGE_STRUCTURES", "-DRPC_STRUCTURES", "-include", "{}/src/common/sdk/nvidia/inc/nvtypes.h", "-I{}/src/nvidia/generated",
"-I{}/src/common/inc", "-I{}/src/nvidia/inc", "-I{}/src/nvidia/interface/", "-I{}/src/nvidia/inc/kernel", "-I{}/src/nvidia/inc/libraries",
"-I{}/src/nvidia/arch/nvalloc/common/inc", "-I{}/kernel-open/nvidia-uvm", "-I{}/kernel-open/common/inc", "-I{}/src/common/sdk/nvidia/inc",
"-I{}/src/nvidia/arch/nvalloc/unix/include", "-I{}/src/common/sdk/nvidia/inc/ctrl"
], tarball=nv_src["nv_570"], anon_names={
"{}/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h:8": "rpc_fns",
"{}/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h:244": "rpc_events"
})
# this defines all syscall numbers. should probably unify linux autogen?
case "io_uring": return load("io_uring", [], ["/usr/include/liburing.h", "/usr/include/linux/io_uring.h", "/usr/include/asm-generic/unistd.h"],
rules=[('__NR', 'NR')])
case "ib": return load("ib", ["ibverbs"], ["/usr/include/infiniband/verbs.h", "/usr/include/infiniband/verbs_api.h",
"/usr/include/infiniband/ib_user_ioctl_verbs.h","/usr/include/rdma/ib_user_verbs.h"], use_errno=True)
case "llvm": return load("llvm", ["LLVM_PATH"], lambda: [system("llvm-config-20 --includedir")+"/llvm-c/**/*.h"],
args=lambda: system("llvm-config-20 --cflags").split(), recsym=True,
prolog=["from tinygrad.runtime.support.llvm import LLVM_PATH"])
case "pci": return load("pci", [], ["/usr/include/linux/pci_regs.h"])
case "vfio": return load("vfio", [], ["/usr/include/linux/vfio.h"])
# could add rule: WGPU_COMMA -> ','
case "webgpu":
return load("webgpu", ["WEBGPU_PATH"], [root/"extra/webgpu/webgpu.h"], prolog=["from tinygrad.runtime.support.webgpu import WEBGPU_PATH"])
case "libusb": return load("libusb", ["os.getenv('LIBUSB_PATH', find_library('usb-1.0'))"], ["/usr/include/libusb-1.0/libusb.h"])
case "hip": return load("hip", ["os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamdhip64.so'"], ["/opt/rocm/include/hip/hip_ext.h",
"/opt/rocm/include/hip/hiprtc.h", "/opt/rocm/include/hip/hip_runtime_api.h", "/opt/rocm/include/hip/driver_types.h"],
args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"])
case "comgr" | "comgr_3":
return load("comgr_3" if nm == "comgr_3" else "comgr", [
"os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so'", "'/usr/local/lib/libamd_comgr.dylib'", "'/opt/homebrew/lib/libamd_comgr.dylib'"
], ["/opt/rocm/include/amd_comgr/amd_comgr.h"], args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"])
case "hsa": return load("hsa", ["os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libhsa-runtime64.so'", "find_library('hsa-runtime64')"], [
f"/opt/rocm/include/hsa/{s}.h" for s in ["hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue", "amd_hsa_kernel_code", "hsa_ext_finalize",
"hsa_ext_image", "hsa_ven_amd_aqlprofile"] ], args=["-I/opt/rocm/include"])
case "amd_gpu": return load("amd_gpu", [], [root/f"extra/hip_gpu_driver/{s}.h" for s in ["sdma_registers", "nvd", "gc_11_0_0_offset",
"sienna_cichlid_ip_offset"]],
args=["-I/opt/rocm/include", "-x", "c++"])
case "kgsl": return load("kgsl", [], [root/"extra/qcom_gpu_driver/msm_kgsl.h"], args=["-D__user="])
case "adreno": return load("adreno", [], [root/"extra/qcom_gpu_driver/a6xx.xml.h"])
case "qcom_dsp":
return load("qcom_dsp", [], [root/f"extra/dsp/include/{s}.h" for s in ["ion", "msm_ion", "adsprpc_shared", "remote_default", "apps_std"]])
case "sqtt": return load("sqtt", [], [root/"extra/sqtt/sqtt.h"])
case "rocprof":
return load("rocprof", ["find_library('rocprof-trace-decoder')", p:="'/usr/local/lib/rocprof-trace-decoder.so'", p.replace('so','dylib')],
[f"{{}}/include/{s}.h" for s in ["rocprof_trace_decoder", "trace_decoder_instrument", "trace_decoder_types"]],
tarball="https://github.com/ROCm/rocprof-trace-decoder/archive/dd0485100971522cc4cd8ae136bdda431061a04d.tar.gz")
case "mesa": return load("mesa", ["find_library('tinymesa_cpu')",
"(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if OSX else '.so')",
"f'{BASE}/libtinymesa{EXT}'", "'/opt/homebrew/lib/libtinymesa_cpu.dylib'", "'/opt/homebrew/lib/libtinymesa.dylib'"], [
*[f"{{}}/src/compiler/nir/{s}.h" for s in ["nir", "nir_builder", "nir_shader_compiler_options", "nir_serialize"]], "{}/gen/nir_intrinsics.h",
*[f"{{}}/src/nouveau/{s}.h" for s in ["headers/nv_device_info", "compiler/nak"]],
*[f"{{}}/src/gallium/auxiliary/gallivm/lp_bld{s}.h" for s in ["", "_passmgr", "_misc", "_type", "_init", "_nir", "_struct", "_jit_types",
"_flow", "_const"]],
"{}/src/compiler/glsl_types.h", "{}/src/util/blob.h", "{}/src/util/ralloc.h"], args=lambda:[
"-DHAVE_ENDIAN_H", "-DHAVE_STRUCT_TIMESPEC", "-DHAVE_PTHREAD", "-DHAVE_FUNC_ATTRIBUTE_PACKED", "-I{}/src", "-I{}/include", "-I{}/gen",
"-I{}/src/compiler/nir", "-I{}/src/gallium/auxiliary", "-I{}/src/gallium/include", f"-I{system('llvm-config-20 --includedir')}"],
preprocess=lambda path: subprocess.run("""mkdir -p gen/util/format
python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
python3 src/compiler/builtin_types_h.py gen/builtin_types.h""", cwd=path, shell=True, check=True),
tarball="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.4/mesa-25.2.4.tar.gz",
prolog=["import gzip, base64", "from tinygrad.helpers import OSX"], epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
case _: raise AttributeError(f"no such autogen: {nm}")
tinygrad/runtime/autogen/adreno.py
+7806 -17903
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/am/__init__.py
+23
View File
@@ -0,0 +1,23 @@
from tinygrad.runtime.autogen import load, root
am_src="https://github.com/ROCm/ROCK-Kernel-Driver/archive/ceb12c04e2b5b53ec0779362831f5ee40c4921e4.tar.gz"
AMD="{}/drivers/gpu/drm/amd"
inc = ["-include", "stdint.h"]
def __getattr__(nm):
match nm:
case "am": return load("am/am", [], [root/f"extra/amdpci/headers/{s}.h" for s in ["v11_structs", "v12_structs", "amdgpu_vm", "discovery",
"amdgpu_ucode", "psp_gfx_if", "amdgpu_psp", "amdgpu_irq", "amdgpu_doorbell"]]+[f"{AMD}/include/soc15_ih_clientid.h"], args=inc, tarball=am_src)
case "pm4_soc15": return load("am/pm4_soc15", [], [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/soc15d.h"], tarball=am_src)
case "pm4_nv": return load("am/pm4_nv", [], [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/nvd.h"], tarball=am_src)
case "sdma_4_0_0": return load("am/sdma_4_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/vega10_sdma_pkt_open.h"],
args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
case "sdma_5_0_0": return load("am/sdma_5_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/navi10_sdma_pkt_open.h"],
args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
case "sdma_6_0_0": return load("am/sdma_6_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}//amdgpu/sdma_v6_0_0_pkt_open.h"],
args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
case "smu_v13_0_0": return load("am/smu_v13_0_0",[],[f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_0_ppsmc","smu13_driver_if_v13_0_0"]]
+[root/"extra/amdpci/headers/amdgpu_smu.h"], tarball=am_src),
case "smu_v14_0_2": return load("am/smu_v14_0_2", [], [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v14_0_0_pmfw", "smu_v14_0_2_ppsmc",
"smu14_driver_if_v14_0"]]+[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, tarball=am_src)
case _: raise AttributeError(f"no such autogen: {nm}")
tinygrad/runtime/autogen/am/am.py
+3899 -5626
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/am/pm4_nv.py
+436 -939
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/am/pm4_soc15.py
+424 -908
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/am/sdma_4_0_0.py
+2637 -5209
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/am/sdma_4_4_2.py
-5211
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/am/sdma_5_0_0.py
+3589 -7103
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/am/sdma_6_0_0.py
+4085 -8085
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/am/smu_v13_0_0.py
+990 -3030
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/am/smu_v14_0_2.py
+1871 -3494
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/amd_gpu.py
+13009 -21999
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/comgr.py
+393 -911
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/comgr_3.py
+387 -894
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/cuda.py
+3580 -7337
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/hip.py
+3328 -5890
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/hsa.py
+2354 -5627
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/ib.py
+3496 -7170
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/io_uring.py
+960 -2291
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/kfd.py
+628 -1399
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/kgsl.py
+595 -1242
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/libusb.py
+774 -1589
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/llvm.py
+13685 -11232
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/mesa.py
+8436 -19660
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/nv.py
+4837
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/nv/nv.py
-8740
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/nv_570.py
+20570
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/nv_580.py
+21640
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/nv_gpu.py
-38764
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/nvjitlink.py
+89
View File
@@ -0,0 +1,89 @@
# mypy: ignore-errors
import ctypes
from tinygrad.helpers import unwrap
from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
from ctypes.util import find_library
def dll():
try: return ctypes.CDLL(unwrap(find_library('nvJitLink')))
except: pass
return None
dll = dll()
nvJitLinkResult = CEnum(ctypes.c_uint32)
NVJITLINK_SUCCESS = nvJitLinkResult.define('NVJITLINK_SUCCESS', 0)
NVJITLINK_ERROR_UNRECOGNIZED_OPTION = nvJitLinkResult.define('NVJITLINK_ERROR_UNRECOGNIZED_OPTION', 1)
NVJITLINK_ERROR_MISSING_ARCH = nvJitLinkResult.define('NVJITLINK_ERROR_MISSING_ARCH', 2)
NVJITLINK_ERROR_INVALID_INPUT = nvJitLinkResult.define('NVJITLINK_ERROR_INVALID_INPUT', 3)
NVJITLINK_ERROR_PTX_COMPILE = nvJitLinkResult.define('NVJITLINK_ERROR_PTX_COMPILE', 4)
NVJITLINK_ERROR_NVVM_COMPILE = nvJitLinkResult.define('NVJITLINK_ERROR_NVVM_COMPILE', 5)
NVJITLINK_ERROR_INTERNAL = nvJitLinkResult.define('NVJITLINK_ERROR_INTERNAL', 6)
nvJitLinkInputType = CEnum(ctypes.c_uint32)
NVJITLINK_INPUT_NONE = nvJitLinkInputType.define('NVJITLINK_INPUT_NONE', 0)
NVJITLINK_INPUT_CUBIN = nvJitLinkInputType.define('NVJITLINK_INPUT_CUBIN', 1)
NVJITLINK_INPUT_PTX = nvJitLinkInputType.define('NVJITLINK_INPUT_PTX', 2)
NVJITLINK_INPUT_LTOIR = nvJitLinkInputType.define('NVJITLINK_INPUT_LTOIR', 3)
NVJITLINK_INPUT_FATBIN = nvJitLinkInputType.define('NVJITLINK_INPUT_FATBIN', 4)
NVJITLINK_INPUT_OBJECT = nvJitLinkInputType.define('NVJITLINK_INPUT_OBJECT', 5)
NVJITLINK_INPUT_LIBRARY = nvJitLinkInputType.define('NVJITLINK_INPUT_LIBRARY', 6)
class struct_nvJitLink(Struct): pass
nvJitLinkHandle = ctypes.POINTER(struct_nvJitLink)
uint32_t = ctypes.c_uint32
# nvJitLinkResult nvJitLinkCreate(nvJitLinkHandle *handle, uint32_t numOptions, const char **options)
try: (nvJitLinkCreate:=dll.nvJitLinkCreate).restype, nvJitLinkCreate.argtypes = nvJitLinkResult, [ctypes.POINTER(nvJitLinkHandle), uint32_t, ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
except AttributeError: pass
# nvJitLinkResult nvJitLinkDestroy(nvJitLinkHandle *handle)
try: (nvJitLinkDestroy:=dll.nvJitLinkDestroy).restype, nvJitLinkDestroy.argtypes = nvJitLinkResult, [ctypes.POINTER(nvJitLinkHandle)]
except AttributeError: pass
size_t = ctypes.c_uint64
# nvJitLinkResult nvJitLinkAddData(nvJitLinkHandle handle, nvJitLinkInputType inputType, const void *data, size_t size, const char *name)
try: (nvJitLinkAddData:=dll.nvJitLinkAddData).restype, nvJitLinkAddData.argtypes = nvJitLinkResult, [nvJitLinkHandle, nvJitLinkInputType, ctypes.c_void_p, size_t, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvJitLinkResult nvJitLinkAddFile(nvJitLinkHandle handle, nvJitLinkInputType inputType, const char *fileName)
try: (nvJitLinkAddFile:=dll.nvJitLinkAddFile).restype, nvJitLinkAddFile.argtypes = nvJitLinkResult, [nvJitLinkHandle, nvJitLinkInputType, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvJitLinkResult nvJitLinkComplete(nvJitLinkHandle handle)
try: (nvJitLinkComplete:=dll.nvJitLinkComplete).restype, nvJitLinkComplete.argtypes = nvJitLinkResult, [nvJitLinkHandle]
except AttributeError: pass
# nvJitLinkResult nvJitLinkGetLinkedCubinSize(nvJitLinkHandle handle, size_t *size)
try: (nvJitLinkGetLinkedCubinSize:=dll.nvJitLinkGetLinkedCubinSize).restype, nvJitLinkGetLinkedCubinSize.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(size_t)]
except AttributeError: pass
# nvJitLinkResult nvJitLinkGetLinkedCubin(nvJitLinkHandle handle, void *cubin)
try: (nvJitLinkGetLinkedCubin:=dll.nvJitLinkGetLinkedCubin).restype, nvJitLinkGetLinkedCubin.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.c_void_p]
except AttributeError: pass
# nvJitLinkResult nvJitLinkGetLinkedPtxSize(nvJitLinkHandle handle, size_t *size)
try: (nvJitLinkGetLinkedPtxSize:=dll.nvJitLinkGetLinkedPtxSize).restype, nvJitLinkGetLinkedPtxSize.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(size_t)]
except AttributeError: pass
# nvJitLinkResult nvJitLinkGetLinkedPtx(nvJitLinkHandle handle, char *ptx)
try: (nvJitLinkGetLinkedPtx:=dll.nvJitLinkGetLinkedPtx).restype, nvJitLinkGetLinkedPtx.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvJitLinkResult nvJitLinkGetErrorLogSize(nvJitLinkHandle handle, size_t *size)
try: (nvJitLinkGetErrorLogSize:=dll.nvJitLinkGetErrorLogSize).restype, nvJitLinkGetErrorLogSize.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(size_t)]
except AttributeError: pass
# nvJitLinkResult nvJitLinkGetErrorLog(nvJitLinkHandle handle, char *log)
try: (nvJitLinkGetErrorLog:=dll.nvJitLinkGetErrorLog).restype, nvJitLinkGetErrorLog.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvJitLinkResult nvJitLinkGetInfoLogSize(nvJitLinkHandle handle, size_t *size)
try: (nvJitLinkGetInfoLogSize:=dll.nvJitLinkGetInfoLogSize).restype, nvJitLinkGetInfoLogSize.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(size_t)]
except AttributeError: pass
# nvJitLinkResult nvJitLinkGetInfoLog(nvJitLinkHandle handle, char *log)
try: (nvJitLinkGetInfoLog:=dll.nvJitLinkGetInfoLog).restype, nvJitLinkGetInfoLog.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvJitLinkResult nvJitLinkVersion(unsigned int *major, unsigned int *minor)
try: (nvJitLinkVersion:=dll.nvJitLinkVersion).restype, nvJitLinkVersion.argtypes = nvJitLinkResult, [ctypes.POINTER(ctypes.c_uint32), ctypes.POINTER(ctypes.c_uint32)]
except AttributeError: pass
tinygrad/runtime/autogen/nvrtc.py
+101 -569
View File
@@ -1,581 +1,113 @@
# mypy: ignore-errors
# -*- coding: utf-8 -*-
#
# TARGET arch is: []
# WORD_SIZE is: 8
# POINTER_SIZE is: 8
# LONGDOUBLE_SIZE is: 16
#
import ctypes, ctypes.util
import ctypes
from tinygrad.helpers import unwrap
from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
from ctypes.util import find_library
def dll():
try: return ctypes.CDLL(unwrap(find_library('nvrtc')))
except: pass
return None
dll = dll()
nvrtcResult = CEnum(ctypes.c_uint32)
NVRTC_SUCCESS = nvrtcResult.define('NVRTC_SUCCESS', 0)
NVRTC_ERROR_OUT_OF_MEMORY = nvrtcResult.define('NVRTC_ERROR_OUT_OF_MEMORY', 1)
NVRTC_ERROR_PROGRAM_CREATION_FAILURE = nvrtcResult.define('NVRTC_ERROR_PROGRAM_CREATION_FAILURE', 2)
NVRTC_ERROR_INVALID_INPUT = nvrtcResult.define('NVRTC_ERROR_INVALID_INPUT', 3)
NVRTC_ERROR_INVALID_PROGRAM = nvrtcResult.define('NVRTC_ERROR_INVALID_PROGRAM', 4)
NVRTC_ERROR_INVALID_OPTION = nvrtcResult.define('NVRTC_ERROR_INVALID_OPTION', 5)
NVRTC_ERROR_COMPILATION = nvrtcResult.define('NVRTC_ERROR_COMPILATION', 6)
NVRTC_ERROR_BUILTIN_OPERATION_FAILURE = nvrtcResult.define('NVRTC_ERROR_BUILTIN_OPERATION_FAILURE', 7)
NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION = nvrtcResult.define('NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION', 8)
NVRTC_ERROR_NO_LOWERED_NAMES_BEFORE_COMPILATION = nvrtcResult.define('NVRTC_ERROR_NO_LOWERED_NAMES_BEFORE_COMPILATION', 9)
NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID = nvrtcResult.define('NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID', 10)
NVRTC_ERROR_INTERNAL_ERROR = nvrtcResult.define('NVRTC_ERROR_INTERNAL_ERROR', 11)
_libraries = {}
_libraries['libnvrtc.so'] = ctypes.CDLL(ctypes.util.find_library('nvrtc'))
def string_cast(char_pointer, encoding='utf-8', errors='strict'):
value = ctypes.cast(char_pointer, ctypes.c_char_p).value
if value is not None and encoding is not None:
value = value.decode(encoding, errors=errors)
return value
# const char *nvrtcGetErrorString(nvrtcResult result)
try: (nvrtcGetErrorString:=dll.nvrtcGetErrorString).restype, nvrtcGetErrorString.argtypes = ctypes.POINTER(ctypes.c_char), [nvrtcResult]
except AttributeError: pass
# nvrtcResult nvrtcVersion(int *major, int *minor)
try: (nvrtcVersion:=dll.nvrtcVersion).restype, nvrtcVersion.argtypes = nvrtcResult, [ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32)]
except AttributeError: pass
def char_pointer_cast(string, encoding='utf-8'):
if encoding is not None:
try:
string = string.encode(encoding)
except AttributeError:
# In Python3, bytes has no encode attribute
pass
string = ctypes.c_char_p(string)
return ctypes.cast(string, ctypes.POINTER(ctypes.c_char))
# nvrtcResult nvrtcGetNumSupportedArchs(int *numArchs)
try: (nvrtcGetNumSupportedArchs:=dll.nvrtcGetNumSupportedArchs).restype, nvrtcGetNumSupportedArchs.argtypes = nvrtcResult, [ctypes.POINTER(ctypes.c_int32)]
except AttributeError: pass
# nvrtcResult nvrtcGetSupportedArchs(int *supportedArchs)
try: (nvrtcGetSupportedArchs:=dll.nvrtcGetSupportedArchs).restype, nvrtcGetSupportedArchs.argtypes = nvrtcResult, [ctypes.POINTER(ctypes.c_int32)]
except AttributeError: pass
class AsDictMixin:
import sys
if sys.version_info >= (3, 14): _layout_ = 'ms'
@classmethod
def as_dict(cls, self):
result = {}
if not isinstance(self, AsDictMixin):
# not a structure, assume it's already a python object
return self
if not hasattr(cls, "_fields_"):
return result
# sys.version_info >= (3, 5)
# for (field, *_) in cls._fields_: # noqa
for field_tuple in cls._fields_: # noqa
field = field_tuple[0]
if field.startswith('PADDING_'):
continue
value = getattr(self, field)
type_ = type(value)
if hasattr(value, "_length_") and hasattr(value, "_type_"):
# array
if not hasattr(type_, "as_dict"):
value = [v for v in value]
else:
type_ = type_._type_
value = [type_.as_dict(v) for v in value]
elif hasattr(value, "contents") and hasattr(value, "_type_"):
# pointer
try:
if not hasattr(type_, "as_dict"):
value = value.contents
else:
type_ = type_._type_
value = type_.as_dict(value.contents)
except ValueError:
# nullptr
value = None
elif isinstance(value, AsDictMixin):
# other structure
value = type_.as_dict(value)
result[field] = value
return result
class Structure(ctypes.Structure, AsDictMixin):
def __init__(self, *args, **kwds):
# We don't want to use positional arguments fill PADDING_* fields
args = dict(zip(self.__class__._field_names_(), args))
args.update(kwds)
super(Structure, self).__init__(**args)
@classmethod
def _field_names_(cls):
if hasattr(cls, '_fields_'):
return (f[0] for f in cls._fields_ if not f[0].startswith('PADDING'))
else:
return ()
@classmethod
def get_type(cls, field):
for f in cls._fields_:
if f[0] == field:
return f[1]
return None
@classmethod
def bind(cls, bound_fields):
fields = {}
for name, type_ in cls._fields_:
if hasattr(type_, "restype"):
if name in bound_fields:
if bound_fields[name] is None:
fields[name] = type_()
else:
# use a closure to capture the callback from the loop scope
fields[name] = (
type_((lambda callback: lambda *args: callback(*args))(
bound_fields[name]))
)
del bound_fields[name]
else:
# default callback implementation (does nothing)
try:
default_ = type_(0).restype().value
except TypeError:
default_ = None
fields[name] = type_((
lambda default_: lambda *args: default_)(default_))
else:
# not a callback function, use default initialization
if name in bound_fields:
fields[name] = bound_fields[name]
del bound_fields[name]
else:
fields[name] = type_()
if len(bound_fields) != 0:
raise ValueError(
"Cannot bind the following unknown callback(s) {}.{}".format(
cls.__name__, bound_fields.keys()
))
return cls(**fields)
class Union(ctypes.Union, AsDictMixin):
pass
_libraries['libnvJitLink.so'] = ctypes.CDLL(ctypes.util.find_library('nvJitLink'))
c_int128 = ctypes.c_ubyte*16
c_uint128 = c_int128
void = None
if ctypes.sizeof(ctypes.c_longdouble) == 16:
c_long_double_t = ctypes.c_longdouble
else:
c_long_double_t = ctypes.c_ubyte*16
# values for enumeration 'c__EA_nvrtcResult'
c__EA_nvrtcResult__enumvalues = {
0: 'NVRTC_SUCCESS',
1: 'NVRTC_ERROR_OUT_OF_MEMORY',
2: 'NVRTC_ERROR_PROGRAM_CREATION_FAILURE',
3: 'NVRTC_ERROR_INVALID_INPUT',
4: 'NVRTC_ERROR_INVALID_PROGRAM',
5: 'NVRTC_ERROR_INVALID_OPTION',
6: 'NVRTC_ERROR_COMPILATION',
7: 'NVRTC_ERROR_BUILTIN_OPERATION_FAILURE',
8: 'NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION',
9: 'NVRTC_ERROR_NO_LOWERED_NAMES_BEFORE_COMPILATION',
10: 'NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID',
11: 'NVRTC_ERROR_INTERNAL_ERROR',
12: 'NVRTC_ERROR_TIME_FILE_WRITE_FAILED',
}
NVRTC_SUCCESS = 0
NVRTC_ERROR_OUT_OF_MEMORY = 1
NVRTC_ERROR_PROGRAM_CREATION_FAILURE = 2
NVRTC_ERROR_INVALID_INPUT = 3
NVRTC_ERROR_INVALID_PROGRAM = 4
NVRTC_ERROR_INVALID_OPTION = 5
NVRTC_ERROR_COMPILATION = 6
NVRTC_ERROR_BUILTIN_OPERATION_FAILURE = 7
NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION = 8
NVRTC_ERROR_NO_LOWERED_NAMES_BEFORE_COMPILATION = 9
NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID = 10
NVRTC_ERROR_INTERNAL_ERROR = 11
NVRTC_ERROR_TIME_FILE_WRITE_FAILED = 12
c__EA_nvrtcResult = ctypes.c_uint32 # enum
nvrtcResult = c__EA_nvrtcResult
nvrtcResult__enumvalues = c__EA_nvrtcResult__enumvalues
try:
nvrtcGetErrorString = _libraries['libnvrtc.so'].nvrtcGetErrorString
nvrtcGetErrorString.restype = ctypes.POINTER(ctypes.c_char)
nvrtcGetErrorString.argtypes = [nvrtcResult]
except AttributeError:
pass
try:
nvrtcVersion = _libraries['libnvrtc.so'].nvrtcVersion
nvrtcVersion.restype = nvrtcResult
nvrtcVersion.argtypes = [ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32)]
except AttributeError:
pass
try:
nvrtcGetNumSupportedArchs = _libraries['libnvrtc.so'].nvrtcGetNumSupportedArchs
nvrtcGetNumSupportedArchs.restype = nvrtcResult
nvrtcGetNumSupportedArchs.argtypes = [ctypes.POINTER(ctypes.c_int32)]
except AttributeError:
pass
try:
nvrtcGetSupportedArchs = _libraries['libnvrtc.so'].nvrtcGetSupportedArchs
nvrtcGetSupportedArchs.restype = nvrtcResult
nvrtcGetSupportedArchs.argtypes = [ctypes.POINTER(ctypes.c_int32)]
except AttributeError:
pass
class struct__nvrtcProgram(Structure):
pass
class struct__nvrtcProgram(Struct): pass
nvrtcProgram = ctypes.POINTER(struct__nvrtcProgram)
try:
nvrtcCreateProgram = _libraries['libnvrtc.so'].nvrtcCreateProgram
nvrtcCreateProgram.restype = nvrtcResult
nvrtcCreateProgram.argtypes = [ctypes.POINTER(ctypes.POINTER(struct__nvrtcProgram)), ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), ctypes.c_int32, ctypes.POINTER(ctypes.POINTER(ctypes.c_char)), ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
except AttributeError:
pass
try:
nvrtcDestroyProgram = _libraries['libnvrtc.so'].nvrtcDestroyProgram
nvrtcDestroyProgram.restype = nvrtcResult
nvrtcDestroyProgram.argtypes = [ctypes.POINTER(ctypes.POINTER(struct__nvrtcProgram))]
except AttributeError:
pass
try:
nvrtcCompileProgram = _libraries['libnvrtc.so'].nvrtcCompileProgram
nvrtcCompileProgram.restype = nvrtcResult
nvrtcCompileProgram.argtypes = [nvrtcProgram, ctypes.c_int32, ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
except AttributeError:
pass
try:
nvrtcGetPTXSize = _libraries['libnvrtc.so'].nvrtcGetPTXSize
nvrtcGetPTXSize.restype = nvrtcResult
nvrtcGetPTXSize.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
nvrtcGetPTX = _libraries['libnvrtc.so'].nvrtcGetPTX
nvrtcGetPTX.restype = nvrtcResult
nvrtcGetPTX.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvrtcGetCUBINSize = _libraries['libnvrtc.so'].nvrtcGetCUBINSize
nvrtcGetCUBINSize.restype = nvrtcResult
nvrtcGetCUBINSize.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
nvrtcGetCUBIN = _libraries['libnvrtc.so'].nvrtcGetCUBIN
nvrtcGetCUBIN.restype = nvrtcResult
nvrtcGetCUBIN.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvrtcGetNVVMSize = _libraries['libnvrtc.so'].nvrtcGetNVVMSize
nvrtcGetNVVMSize.restype = nvrtcResult
nvrtcGetNVVMSize.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
nvrtcGetNVVM = _libraries['libnvrtc.so'].nvrtcGetNVVM
nvrtcGetNVVM.restype = nvrtcResult
nvrtcGetNVVM.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvrtcGetLTOIRSize = _libraries['libnvrtc.so'].nvrtcGetLTOIRSize
nvrtcGetLTOIRSize.restype = nvrtcResult
nvrtcGetLTOIRSize.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
nvrtcGetLTOIR = _libraries['libnvrtc.so'].nvrtcGetLTOIR
nvrtcGetLTOIR.restype = nvrtcResult
nvrtcGetLTOIR.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvrtcGetOptiXIRSize = _libraries['libnvrtc.so'].nvrtcGetOptiXIRSize
nvrtcGetOptiXIRSize.restype = nvrtcResult
nvrtcGetOptiXIRSize.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
nvrtcGetOptiXIR = _libraries['libnvrtc.so'].nvrtcGetOptiXIR
nvrtcGetOptiXIR.restype = nvrtcResult
nvrtcGetOptiXIR.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvrtcGetProgramLogSize = _libraries['libnvrtc.so'].nvrtcGetProgramLogSize
nvrtcGetProgramLogSize.restype = nvrtcResult
nvrtcGetProgramLogSize.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
nvrtcGetProgramLog = _libraries['libnvrtc.so'].nvrtcGetProgramLog
nvrtcGetProgramLog.restype = nvrtcResult
nvrtcGetProgramLog.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvrtcAddNameExpression = _libraries['libnvrtc.so'].nvrtcAddNameExpression
nvrtcAddNameExpression.restype = nvrtcResult
nvrtcAddNameExpression.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvrtcGetLoweredName = _libraries['libnvrtc.so'].nvrtcGetLoweredName
nvrtcGetLoweredName.restype = nvrtcResult
nvrtcGetLoweredName.argtypes = [nvrtcProgram, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
except AttributeError:
pass
# nvrtcResult nvrtcCreateProgram(nvrtcProgram *prog, const char *src, const char *name, int numHeaders, const char *const *headers, const char *const *includeNames)
try: (nvrtcCreateProgram:=dll.nvrtcCreateProgram).restype, nvrtcCreateProgram.argtypes = nvrtcResult, [ctypes.POINTER(nvrtcProgram), ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), ctypes.c_int32, ctypes.POINTER(ctypes.POINTER(ctypes.c_char)), ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
except AttributeError: pass
# values for enumeration 'c__EA_nvJitLinkResult'
c__EA_nvJitLinkResult__enumvalues = {
0: 'NVJITLINK_SUCCESS',
1: 'NVJITLINK_ERROR_UNRECOGNIZED_OPTION',
2: 'NVJITLINK_ERROR_MISSING_ARCH',
3: 'NVJITLINK_ERROR_INVALID_INPUT',
4: 'NVJITLINK_ERROR_PTX_COMPILE',
5: 'NVJITLINK_ERROR_NVVM_COMPILE',
6: 'NVJITLINK_ERROR_INTERNAL',
7: 'NVJITLINK_ERROR_THREADPOOL',
8: 'NVJITLINK_ERROR_UNRECOGNIZED_INPUT',
}
NVJITLINK_SUCCESS = 0
NVJITLINK_ERROR_UNRECOGNIZED_OPTION = 1
NVJITLINK_ERROR_MISSING_ARCH = 2
NVJITLINK_ERROR_INVALID_INPUT = 3
NVJITLINK_ERROR_PTX_COMPILE = 4
NVJITLINK_ERROR_NVVM_COMPILE = 5
NVJITLINK_ERROR_INTERNAL = 6
NVJITLINK_ERROR_THREADPOOL = 7
NVJITLINK_ERROR_UNRECOGNIZED_INPUT = 8
c__EA_nvJitLinkResult = ctypes.c_uint32 # enum
nvJitLinkResult = c__EA_nvJitLinkResult
nvJitLinkResult__enumvalues = c__EA_nvJitLinkResult__enumvalues
# nvrtcResult nvrtcDestroyProgram(nvrtcProgram *prog)
try: (nvrtcDestroyProgram:=dll.nvrtcDestroyProgram).restype, nvrtcDestroyProgram.argtypes = nvrtcResult, [ctypes.POINTER(nvrtcProgram)]
except AttributeError: pass
# values for enumeration 'c__EA_nvJitLinkInputType'
c__EA_nvJitLinkInputType__enumvalues = {
0: 'NVJITLINK_INPUT_NONE',
1: 'NVJITLINK_INPUT_CUBIN',
2: 'NVJITLINK_INPUT_PTX',
3: 'NVJITLINK_INPUT_LTOIR',
4: 'NVJITLINK_INPUT_FATBIN',
5: 'NVJITLINK_INPUT_OBJECT',
6: 'NVJITLINK_INPUT_LIBRARY',
10: 'NVJITLINK_INPUT_ANY',
}
NVJITLINK_INPUT_NONE = 0
NVJITLINK_INPUT_CUBIN = 1
NVJITLINK_INPUT_PTX = 2
NVJITLINK_INPUT_LTOIR = 3
NVJITLINK_INPUT_FATBIN = 4
NVJITLINK_INPUT_OBJECT = 5
NVJITLINK_INPUT_LIBRARY = 6
NVJITLINK_INPUT_ANY = 10
c__EA_nvJitLinkInputType = ctypes.c_uint32 # enum
nvJitLinkInputType = c__EA_nvJitLinkInputType
nvJitLinkInputType__enumvalues = c__EA_nvJitLinkInputType__enumvalues
class struct_nvJitLink(Structure):
pass
# nvrtcResult nvrtcCompileProgram(nvrtcProgram prog, int numOptions, const char *const *options)
try: (nvrtcCompileProgram:=dll.nvrtcCompileProgram).restype, nvrtcCompileProgram.argtypes = nvrtcResult, [nvrtcProgram, ctypes.c_int32, ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
except AttributeError: pass
nvJitLinkHandle = ctypes.POINTER(struct_nvJitLink)
uint32_t = ctypes.c_uint32
try:
__nvJitLinkCreate_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkCreate_12_4
__nvJitLinkCreate_12_4.restype = nvJitLinkResult
__nvJitLinkCreate_12_4.argtypes = [ctypes.POINTER(ctypes.POINTER(struct_nvJitLink)), uint32_t, ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
except AttributeError:
pass
try:
nvJitLinkCreate = _libraries['libnvJitLink.so'].nvJitLinkCreate
nvJitLinkCreate.restype = nvJitLinkResult
nvJitLinkCreate.argtypes = [ctypes.POINTER(ctypes.POINTER(struct_nvJitLink)), uint32_t, ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
except AttributeError:
pass
try:
__nvJitLinkDestroy_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkDestroy_12_4
__nvJitLinkDestroy_12_4.restype = nvJitLinkResult
__nvJitLinkDestroy_12_4.argtypes = [ctypes.POINTER(ctypes.POINTER(struct_nvJitLink))]
except AttributeError:
pass
try:
nvJitLinkDestroy = _libraries['libnvJitLink.so'].nvJitLinkDestroy
nvJitLinkDestroy.restype = nvJitLinkResult
nvJitLinkDestroy.argtypes = [ctypes.POINTER(ctypes.POINTER(struct_nvJitLink))]
except AttributeError:
pass
size_t = ctypes.c_uint64
try:
__nvJitLinkAddData_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkAddData_12_4
__nvJitLinkAddData_12_4.restype = nvJitLinkResult
__nvJitLinkAddData_12_4.argtypes = [nvJitLinkHandle, nvJitLinkInputType, ctypes.POINTER(None), size_t, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvJitLinkAddData = _libraries['libnvJitLink.so'].nvJitLinkAddData
nvJitLinkAddData.restype = nvJitLinkResult
nvJitLinkAddData.argtypes = [nvJitLinkHandle, nvJitLinkInputType, ctypes.POINTER(None), size_t, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
__nvJitLinkAddFile_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkAddFile_12_4
__nvJitLinkAddFile_12_4.restype = nvJitLinkResult
__nvJitLinkAddFile_12_4.argtypes = [nvJitLinkHandle, nvJitLinkInputType, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvJitLinkAddFile = _libraries['libnvJitLink.so'].nvJitLinkAddFile
nvJitLinkAddFile.restype = nvJitLinkResult
nvJitLinkAddFile.argtypes = [nvJitLinkHandle, nvJitLinkInputType, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
__nvJitLinkComplete_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkComplete_12_4
__nvJitLinkComplete_12_4.restype = nvJitLinkResult
__nvJitLinkComplete_12_4.argtypes = [nvJitLinkHandle]
except AttributeError:
pass
try:
nvJitLinkComplete = _libraries['libnvJitLink.so'].nvJitLinkComplete
nvJitLinkComplete.restype = nvJitLinkResult
nvJitLinkComplete.argtypes = [nvJitLinkHandle]
except AttributeError:
pass
try:
__nvJitLinkGetLinkedCubinSize_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkGetLinkedCubinSize_12_4
__nvJitLinkGetLinkedCubinSize_12_4.restype = nvJitLinkResult
__nvJitLinkGetLinkedCubinSize_12_4.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
nvJitLinkGetLinkedCubinSize = _libraries['libnvJitLink.so'].nvJitLinkGetLinkedCubinSize
nvJitLinkGetLinkedCubinSize.restype = nvJitLinkResult
nvJitLinkGetLinkedCubinSize.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
__nvJitLinkGetLinkedCubin_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkGetLinkedCubin_12_4
__nvJitLinkGetLinkedCubin_12_4.restype = nvJitLinkResult
__nvJitLinkGetLinkedCubin_12_4.argtypes = [nvJitLinkHandle, ctypes.POINTER(None)]
except AttributeError:
pass
try:
nvJitLinkGetLinkedCubin = _libraries['libnvJitLink.so'].nvJitLinkGetLinkedCubin
nvJitLinkGetLinkedCubin.restype = nvJitLinkResult
nvJitLinkGetLinkedCubin.argtypes = [nvJitLinkHandle, ctypes.POINTER(None)]
except AttributeError:
pass
try:
__nvJitLinkGetLinkedPtxSize_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkGetLinkedPtxSize_12_4
__nvJitLinkGetLinkedPtxSize_12_4.restype = nvJitLinkResult
__nvJitLinkGetLinkedPtxSize_12_4.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
nvJitLinkGetLinkedPtxSize = _libraries['libnvJitLink.so'].nvJitLinkGetLinkedPtxSize
nvJitLinkGetLinkedPtxSize.restype = nvJitLinkResult
nvJitLinkGetLinkedPtxSize.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
__nvJitLinkGetLinkedPtx_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkGetLinkedPtx_12_4
__nvJitLinkGetLinkedPtx_12_4.restype = nvJitLinkResult
__nvJitLinkGetLinkedPtx_12_4.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvJitLinkGetLinkedPtx = _libraries['libnvJitLink.so'].nvJitLinkGetLinkedPtx
nvJitLinkGetLinkedPtx.restype = nvJitLinkResult
nvJitLinkGetLinkedPtx.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
__nvJitLinkGetErrorLogSize_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkGetErrorLogSize_12_4
__nvJitLinkGetErrorLogSize_12_4.restype = nvJitLinkResult
__nvJitLinkGetErrorLogSize_12_4.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
nvJitLinkGetErrorLogSize = _libraries['libnvJitLink.so'].nvJitLinkGetErrorLogSize
nvJitLinkGetErrorLogSize.restype = nvJitLinkResult
nvJitLinkGetErrorLogSize.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
__nvJitLinkGetErrorLog_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkGetErrorLog_12_4
__nvJitLinkGetErrorLog_12_4.restype = nvJitLinkResult
__nvJitLinkGetErrorLog_12_4.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvJitLinkGetErrorLog = _libraries['libnvJitLink.so'].nvJitLinkGetErrorLog
nvJitLinkGetErrorLog.restype = nvJitLinkResult
nvJitLinkGetErrorLog.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
__nvJitLinkGetInfoLogSize_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkGetInfoLogSize_12_4
__nvJitLinkGetInfoLogSize_12_4.restype = nvJitLinkResult
__nvJitLinkGetInfoLogSize_12_4.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
nvJitLinkGetInfoLogSize = _libraries['libnvJitLink.so'].nvJitLinkGetInfoLogSize
nvJitLinkGetInfoLogSize.restype = nvJitLinkResult
nvJitLinkGetInfoLogSize.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_uint64)]
except AttributeError:
pass
try:
__nvJitLinkGetInfoLog_12_4 = _libraries['libnvJitLink.so'].__nvJitLinkGetInfoLog_12_4
__nvJitLinkGetInfoLog_12_4.restype = nvJitLinkResult
__nvJitLinkGetInfoLog_12_4.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvJitLinkGetInfoLog = _libraries['libnvJitLink.so'].nvJitLinkGetInfoLog
nvJitLinkGetInfoLog.restype = nvJitLinkResult
nvJitLinkGetInfoLog.argtypes = [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
except AttributeError:
pass
try:
nvJitLinkVersion = _libraries['libnvJitLink.so'].nvJitLinkVersion
nvJitLinkVersion.restype = nvJitLinkResult
nvJitLinkVersion.argtypes = [ctypes.POINTER(ctypes.c_uint32), ctypes.POINTER(ctypes.c_uint32)]
except AttributeError:
pass
__all__ = \
['NVJITLINK_ERROR_INTERNAL', 'NVJITLINK_ERROR_INVALID_INPUT',
'NVJITLINK_ERROR_MISSING_ARCH', 'NVJITLINK_ERROR_NVVM_COMPILE',
'NVJITLINK_ERROR_PTX_COMPILE', 'NVJITLINK_ERROR_THREADPOOL',
'NVJITLINK_ERROR_UNRECOGNIZED_INPUT',
'NVJITLINK_ERROR_UNRECOGNIZED_OPTION', 'NVJITLINK_INPUT_ANY',
'NVJITLINK_INPUT_CUBIN', 'NVJITLINK_INPUT_FATBIN',
'NVJITLINK_INPUT_LIBRARY', 'NVJITLINK_INPUT_LTOIR',
'NVJITLINK_INPUT_NONE', 'NVJITLINK_INPUT_OBJECT',
'NVJITLINK_INPUT_PTX', 'NVJITLINK_SUCCESS',
'NVRTC_ERROR_BUILTIN_OPERATION_FAILURE',
'NVRTC_ERROR_COMPILATION', 'NVRTC_ERROR_INTERNAL_ERROR',
'NVRTC_ERROR_INVALID_INPUT', 'NVRTC_ERROR_INVALID_OPTION',
'NVRTC_ERROR_INVALID_PROGRAM',
'NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID',
'NVRTC_ERROR_NO_LOWERED_NAMES_BEFORE_COMPILATION',
'NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION',
'NVRTC_ERROR_OUT_OF_MEMORY',
'NVRTC_ERROR_PROGRAM_CREATION_FAILURE',
'NVRTC_ERROR_TIME_FILE_WRITE_FAILED', 'NVRTC_SUCCESS',
'__nvJitLinkAddData_12_4', '__nvJitLinkAddFile_12_4',
'__nvJitLinkComplete_12_4', '__nvJitLinkCreate_12_4',
'__nvJitLinkDestroy_12_4', '__nvJitLinkGetErrorLogSize_12_4',
'__nvJitLinkGetErrorLog_12_4', '__nvJitLinkGetInfoLogSize_12_4',
'__nvJitLinkGetInfoLog_12_4',
'__nvJitLinkGetLinkedCubinSize_12_4',
'__nvJitLinkGetLinkedCubin_12_4',
'__nvJitLinkGetLinkedPtxSize_12_4',
'__nvJitLinkGetLinkedPtx_12_4', 'c__EA_nvJitLinkInputType',
'c__EA_nvJitLinkResult', 'c__EA_nvrtcResult', 'nvJitLinkAddData',
'nvJitLinkAddFile', 'nvJitLinkComplete', 'nvJitLinkCreate',
'nvJitLinkDestroy', 'nvJitLinkGetErrorLog',
'nvJitLinkGetErrorLogSize', 'nvJitLinkGetInfoLog',
'nvJitLinkGetInfoLogSize', 'nvJitLinkGetLinkedCubin',
'nvJitLinkGetLinkedCubinSize', 'nvJitLinkGetLinkedPtx',
'nvJitLinkGetLinkedPtxSize', 'nvJitLinkHandle',
'nvJitLinkInputType', 'nvJitLinkInputType__enumvalues',
'nvJitLinkResult', 'nvJitLinkResult__enumvalues',
'nvJitLinkVersion', 'nvrtcAddNameExpression',
'nvrtcCompileProgram', 'nvrtcCreateProgram',
'nvrtcDestroyProgram', 'nvrtcGetCUBIN', 'nvrtcGetCUBINSize',
'nvrtcGetErrorString', 'nvrtcGetLTOIR', 'nvrtcGetLTOIRSize',
'nvrtcGetLoweredName', 'nvrtcGetNVVM', 'nvrtcGetNVVMSize',
'nvrtcGetNumSupportedArchs', 'nvrtcGetOptiXIR',
'nvrtcGetOptiXIRSize', 'nvrtcGetPTX', 'nvrtcGetPTXSize',
'nvrtcGetProgramLog', 'nvrtcGetProgramLogSize',
'nvrtcGetSupportedArchs', 'nvrtcProgram', 'nvrtcResult',
'nvrtcResult__enumvalues', 'nvrtcVersion', 'size_t',
'struct__nvrtcProgram', 'struct_nvJitLink', 'uint32_t']
# nvrtcResult nvrtcGetPTXSize(nvrtcProgram prog, size_t *ptxSizeRet)
try: (nvrtcGetPTXSize:=dll.nvrtcGetPTXSize).restype, nvrtcGetPTXSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
except AttributeError: pass
# nvrtcResult nvrtcGetPTX(nvrtcProgram prog, char *ptx)
try: (nvrtcGetPTX:=dll.nvrtcGetPTX).restype, nvrtcGetPTX.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvrtcResult nvrtcGetCUBINSize(nvrtcProgram prog, size_t *cubinSizeRet)
try: (nvrtcGetCUBINSize:=dll.nvrtcGetCUBINSize).restype, nvrtcGetCUBINSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
except AttributeError: pass
# nvrtcResult nvrtcGetCUBIN(nvrtcProgram prog, char *cubin)
try: (nvrtcGetCUBIN:=dll.nvrtcGetCUBIN).restype, nvrtcGetCUBIN.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# __attribute__((deprecated("This function will be removed in a future release. Please use nvrtcGetLTOIRSize instead"))) nvrtcResult nvrtcGetNVVMSize(nvrtcProgram prog, size_t *nvvmSizeRet)
try: (nvrtcGetNVVMSize:=dll.nvrtcGetNVVMSize).restype, nvrtcGetNVVMSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
except AttributeError: pass
# __attribute__((deprecated("This function will be removed in a future release. Please use nvrtcGetLTOIR instead"))) nvrtcResult nvrtcGetNVVM(nvrtcProgram prog, char *nvvm)
try: (nvrtcGetNVVM:=dll.nvrtcGetNVVM).restype, nvrtcGetNVVM.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvrtcResult nvrtcGetLTOIRSize(nvrtcProgram prog, size_t *LTOIRSizeRet)
try: (nvrtcGetLTOIRSize:=dll.nvrtcGetLTOIRSize).restype, nvrtcGetLTOIRSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
except AttributeError: pass
# nvrtcResult nvrtcGetLTOIR(nvrtcProgram prog, char *LTOIR)
try: (nvrtcGetLTOIR:=dll.nvrtcGetLTOIR).restype, nvrtcGetLTOIR.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvrtcResult nvrtcGetOptiXIRSize(nvrtcProgram prog, size_t *optixirSizeRet)
try: (nvrtcGetOptiXIRSize:=dll.nvrtcGetOptiXIRSize).restype, nvrtcGetOptiXIRSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
except AttributeError: pass
# nvrtcResult nvrtcGetOptiXIR(nvrtcProgram prog, char *optixir)
try: (nvrtcGetOptiXIR:=dll.nvrtcGetOptiXIR).restype, nvrtcGetOptiXIR.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvrtcResult nvrtcGetProgramLogSize(nvrtcProgram prog, size_t *logSizeRet)
try: (nvrtcGetProgramLogSize:=dll.nvrtcGetProgramLogSize).restype, nvrtcGetProgramLogSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
except AttributeError: pass
# nvrtcResult nvrtcGetProgramLog(nvrtcProgram prog, char *log)
try: (nvrtcGetProgramLog:=dll.nvrtcGetProgramLog).restype, nvrtcGetProgramLog.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvrtcResult nvrtcAddNameExpression(nvrtcProgram prog, const char *const name_expression)
try: (nvrtcAddNameExpression:=dll.nvrtcAddNameExpression).restype, nvrtcAddNameExpression.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
except AttributeError: pass
# nvrtcResult nvrtcGetLoweredName(nvrtcProgram prog, const char *const name_expression, const char **lowered_name)
try: (nvrtcGetLoweredName:=dll.nvrtcGetLoweredName).restype, nvrtcGetLoweredName.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
except AttributeError: pass
__DEPRECATED__ = lambda msg: __attribute__((deprecated(msg)))
tinygrad/runtime/autogen/opencl.py
+964 -1716
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/pci.py
+962 -1330
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/qcom_dsp.py
+535 -1652
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/rocprof.py
+234 -643
View File
@@ -1,669 +1,260 @@
# pylint: skip-file
# mypy: ignore-errors
# -*- coding: utf-8 -*-
#
# TARGET arch is: []
# WORD_SIZE is: 8
# POINTER_SIZE is: 8
# LONGDOUBLE_SIZE is: 16
#
import ctypes, ctypes.util
PATHS_TO_TRY = [
'/usr/local/lib/librocprof-trace-decoder.so',
'/usr/local/lib/librocprof-trace-decoder.dylib',
]
def _try_dlopen_rocprof_trace_decoder():
library = ctypes.util.find_library("rocprof-trace-decoder")
if library:
try: return ctypes.CDLL(library)
except OSError: pass
for candidate in PATHS_TO_TRY:
try: return ctypes.CDLL(candidate)
except OSError: pass
import ctypes
from tinygrad.helpers import unwrap
from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
from ctypes.util import find_library
def dll():
try: return ctypes.CDLL(unwrap(find_library('rocprof-trace-decoder')))
except: pass
try: return ctypes.CDLL(unwrap('/usr/local/lib/rocprof-trace-decoder.so'))
except: pass
try: return ctypes.CDLL(unwrap('/usr/local/lib/rocprof-trace-decoder.dylib'))
except: pass
return None
dll = dll()
rocprofiler_thread_trace_decoder_status_t = CEnum(ctypes.c_uint32)
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS = rocprofiler_thread_trace_decoder_status_t.define('ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS', 0)
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR = rocprofiler_thread_trace_decoder_status_t.define('ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR', 1)
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES = rocprofiler_thread_trace_decoder_status_t.define('ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES', 2)
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT = rocprofiler_thread_trace_decoder_status_t.define('ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT', 3)
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA = rocprofiler_thread_trace_decoder_status_t.define('ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA', 4)
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST = rocprofiler_thread_trace_decoder_status_t.define('ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST', 5)
class AsDictMixin:
@classmethod
def as_dict(cls, self):
result = {}
if not isinstance(self, AsDictMixin):
# not a structure, assume it's already a python object
return self
if not hasattr(cls, "_fields_"):
return result
# sys.version_info >= (3, 5)
# for (field, *_) in cls._fields_: # noqa
for field_tuple in cls._fields_: # noqa
field = field_tuple[0]
if field.startswith('PADDING_'):
continue
value = getattr(self, field)
type_ = type(value)
if hasattr(value, "_length_") and hasattr(value, "_type_"):
# array
if not hasattr(type_, "as_dict"):
value = [v for v in value]
else:
type_ = type_._type_
value = [type_.as_dict(v) for v in value]
elif hasattr(value, "contents") and hasattr(value, "_type_"):
# pointer
try:
if not hasattr(type_, "as_dict"):
value = value.contents
else:
type_ = type_._type_
value = type_.as_dict(value.contents)
except ValueError:
# nullptr
value = None
elif isinstance(value, AsDictMixin):
# other structure
value = type_.as_dict(value)
result[field] = value
return result
enum_rocprofiler_thread_trace_decoder_record_type_t = CEnum(ctypes.c_uint32)
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP = enum_rocprofiler_thread_trace_decoder_record_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP', 0)
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY = enum_rocprofiler_thread_trace_decoder_record_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY', 1)
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT = enum_rocprofiler_thread_trace_decoder_record_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT', 2)
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE = enum_rocprofiler_thread_trace_decoder_record_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE', 3)
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO = enum_rocprofiler_thread_trace_decoder_record_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO', 4)
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG = enum_rocprofiler_thread_trace_decoder_record_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG', 5)
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA = enum_rocprofiler_thread_trace_decoder_record_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA', 6)
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME = enum_rocprofiler_thread_trace_decoder_record_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME', 7)
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY = enum_rocprofiler_thread_trace_decoder_record_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY', 8)
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST = enum_rocprofiler_thread_trace_decoder_record_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST', 9)
class Structure(ctypes.Structure, AsDictMixin):
def __init__(self, *args, **kwds):
# We don't want to use positional arguments fill PADDING_* fields
args = dict(zip(self.__class__._field_names_(), args))
args.update(kwds)
super(Structure, self).__init__(**args)
@classmethod
def _field_names_(cls):
if hasattr(cls, '_fields_'):
return (f[0] for f in cls._fields_ if not f[0].startswith('PADDING'))
else:
return ()
@classmethod
def get_type(cls, field):
for f in cls._fields_:
if f[0] == field:
return f[1]
return None
@classmethod
def bind(cls, bound_fields):
fields = {}
for name, type_ in cls._fields_:
if hasattr(type_, "restype"):
if name in bound_fields:
if bound_fields[name] is None:
fields[name] = type_()
else:
# use a closure to capture the callback from the loop scope
fields[name] = (
type_((lambda callback: lambda *args: callback(*args))(
bound_fields[name]))
)
del bound_fields[name]
else:
# default callback implementation (does nothing)
try:
default_ = type_(0).restype().value
except TypeError:
default_ = None
fields[name] = type_((
lambda default_: lambda *args: default_)(default_))
else:
# not a callback function, use default initialization
if name in bound_fields:
fields[name] = bound_fields[name]
del bound_fields[name]
else:
fields[name] = type_()
if len(bound_fields) != 0:
raise ValueError(
"Cannot bind the following unknown callback(s) {}.{}".format(
cls.__name__, bound_fields.keys()
))
return cls(**fields)
class Union(ctypes.Union, AsDictMixin):
pass
c_int128 = ctypes.c_ubyte*16
c_uint128 = c_int128
void = None
if ctypes.sizeof(ctypes.c_longdouble) == 16:
c_long_double_t = ctypes.c_longdouble
else:
c_long_double_t = ctypes.c_ubyte*16
def string_cast(char_pointer, encoding='utf-8', errors='strict'):
value = ctypes.cast(char_pointer, ctypes.c_char_p).value
if value is not None and encoding is not None:
value = value.decode(encoding, errors=errors)
return value
def char_pointer_cast(string, encoding='utf-8'):
if encoding is not None:
try:
string = string.encode(encoding)
except AttributeError:
# In Python3, bytes has no encode attribute
pass
string = ctypes.c_char_p(string)
return ctypes.cast(string, ctypes.POINTER(ctypes.c_char))
class FunctionFactoryStub:
def __getattr__(self, _):
return ctypes.CFUNCTYPE(lambda y:y)
# libraries['FIXME_STUB'] explanation
# As you did not list (-l libraryname.so) a library that exports this function
# This is a non-working stub instead.
# You can either re-run clan2py with -l /path/to/library.so
# Or manually fix this by comment the ctypes.CDLL loading
_libraries = {}
_libraries['FIXME_STUB'] = _try_dlopen_rocprof_trace_decoder() # ctypes.CDLL('FIXME_STUB')
# values for enumeration 'rocprofiler_thread_trace_decoder_info_t'
rocprofiler_thread_trace_decoder_info_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST',
2: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE',
3: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE',
4: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST',
}
ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE = 0
ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST = 1
ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE = 2
ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE = 3
ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST = 4
rocprofiler_thread_trace_decoder_info_t = ctypes.c_uint32 # enum
class struct_rocprofiler_thread_trace_decoder_pc_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_pc_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_pc_t._fields_ = [
('address', ctypes.c_uint64),
('code_object_id', ctypes.c_uint64),
]
rocprofiler_thread_trace_decoder_pc_t = struct_rocprofiler_thread_trace_decoder_pc_t
class struct_rocprofiler_thread_trace_decoder_perfevent_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_perfevent_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_perfevent_t._fields_ = [
('time', ctypes.c_int64),
('events0', ctypes.c_uint16),
('events1', ctypes.c_uint16),
('events2', ctypes.c_uint16),
('events3', ctypes.c_uint16),
('CU', ctypes.c_ubyte),
('bank', ctypes.c_ubyte),
('PADDING_0', ctypes.c_ubyte * 6),
]
rocprofiler_thread_trace_decoder_perfevent_t = struct_rocprofiler_thread_trace_decoder_perfevent_t
class struct_rocprofiler_thread_trace_decoder_occupancy_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_occupancy_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_occupancy_t._fields_ = [
('pc', rocprofiler_thread_trace_decoder_pc_t),
('time', ctypes.c_uint64),
('reserved', ctypes.c_ubyte),
('cu', ctypes.c_ubyte),
('simd', ctypes.c_ubyte),
('wave_id', ctypes.c_ubyte),
('start', ctypes.c_uint32, 1),
('_rsvd', ctypes.c_uint32, 31),
]
rocprofiler_thread_trace_decoder_occupancy_t = struct_rocprofiler_thread_trace_decoder_occupancy_t
# values for enumeration 'rocprofiler_thread_trace_decoder_wstate_type_t'
rocprofiler_thread_trace_decoder_wstate_type_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE',
2: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC',
3: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT',
4: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL',
5: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST',
}
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY = 0
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE = 1
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC = 2
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT = 3
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL = 4
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST = 5
rocprofiler_thread_trace_decoder_wstate_type_t = ctypes.c_uint32 # enum
class struct_rocprofiler_thread_trace_decoder_wave_state_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_wave_state_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_wave_state_t._fields_ = [
('type', ctypes.c_int32),
('duration', ctypes.c_int32),
]
rocprofiler_thread_trace_decoder_wave_state_t = struct_rocprofiler_thread_trace_decoder_wave_state_t
# values for enumeration 'rocprofiler_thread_trace_decoder_inst_category_t'
rocprofiler_thread_trace_decoder_inst_category_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM',
2: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU',
3: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM',
4: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT',
5: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS',
6: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU',
7: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP',
8: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT',
9: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED',
10: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT',
11: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE',
12: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH',
13: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST',
}
ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE = 0
ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM = 1
ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU = 2
ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM = 3
ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT = 4
ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS = 5
ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU = 6
ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP = 7
ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT = 8
ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED = 9
ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT = 10
ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE = 11
ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH = 12
ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST = 13
rocprofiler_thread_trace_decoder_inst_category_t = ctypes.c_uint32 # enum
class struct_rocprofiler_thread_trace_decoder_inst_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_inst_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_inst_t._fields_ = [
('category', ctypes.c_uint32, 8),
('stall', ctypes.c_uint32, 24),
('duration', ctypes.c_int32),
('time', ctypes.c_int64),
('pc', rocprofiler_thread_trace_decoder_pc_t),
]
rocprofiler_thread_trace_decoder_inst_t = struct_rocprofiler_thread_trace_decoder_inst_t
class struct_rocprofiler_thread_trace_decoder_wave_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_wave_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_wave_t._fields_ = [
('cu', ctypes.c_ubyte),
('simd', ctypes.c_ubyte),
('wave_id', ctypes.c_ubyte),
('contexts', ctypes.c_ubyte),
('_rsvd1', ctypes.c_uint32),
('_rsvd2', ctypes.c_uint32),
('_rsvd3', ctypes.c_uint32),
('begin_time', ctypes.c_int64),
('end_time', ctypes.c_int64),
('timeline_size', ctypes.c_uint64),
('instructions_size', ctypes.c_uint64),
('timeline_array', ctypes.POINTER(struct_rocprofiler_thread_trace_decoder_wave_state_t)),
('instructions_array', ctypes.POINTER(struct_rocprofiler_thread_trace_decoder_inst_t)),
]
rocprofiler_thread_trace_decoder_wave_t = struct_rocprofiler_thread_trace_decoder_wave_t
class struct_rocprofiler_thread_trace_decoder_realtime_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_realtime_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_realtime_t._fields_ = [
('shader_clock', ctypes.c_int64),
('realtime_clock', ctypes.c_uint64),
('reserved', ctypes.c_uint64),
]
rocprofiler_thread_trace_decoder_realtime_t = struct_rocprofiler_thread_trace_decoder_realtime_t
# values for enumeration 'rocprofiler_thread_trace_decoder_shaderdata_flags_t'
rocprofiler_thread_trace_decoder_shaderdata_flags_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV',
}
ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM = 0
ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV = 1
rocprofiler_thread_trace_decoder_shaderdata_flags_t = ctypes.c_uint32 # enum
class struct_rocprofiler_thread_trace_decoder_shaderdata_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_shaderdata_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_shaderdata_t._fields_ = [
('time', ctypes.c_int64),
('value', ctypes.c_uint64),
('cu', ctypes.c_ubyte),
('simd', ctypes.c_ubyte),
('wave_id', ctypes.c_ubyte),
('flags', ctypes.c_ubyte),
('reserved', ctypes.c_uint32),
]
rocprofiler_thread_trace_decoder_shaderdata_t = struct_rocprofiler_thread_trace_decoder_shaderdata_t
# values for enumeration 'rocprofiler_thread_trace_decoder_record_type_t'
rocprofiler_thread_trace_decoder_record_type_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY',
2: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT',
3: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE',
4: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO',
5: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG',
6: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA',
7: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME',
8: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY',
9: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST',
}
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP = 0
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY = 1
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT = 2
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE = 3
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO = 4
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG = 5
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA = 6
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME = 7
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY = 8
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST = 9
rocprofiler_thread_trace_decoder_record_type_t = ctypes.c_uint32 # enum
# values for enumeration 'c__EA_rocprofiler_thread_trace_decoder_status_t'
c__EA_rocprofiler_thread_trace_decoder_status_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR',
2: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES',
3: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT',
4: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA',
5: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST',
}
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS = 0
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR = 1
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES = 2
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT = 3
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA = 4
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST = 5
c__EA_rocprofiler_thread_trace_decoder_status_t = ctypes.c_uint32 # enum
rocprofiler_thread_trace_decoder_status_t = c__EA_rocprofiler_thread_trace_decoder_status_t
rocprofiler_thread_trace_decoder_status_t__enumvalues = c__EA_rocprofiler_thread_trace_decoder_status_t__enumvalues
rocprof_trace_decoder_trace_callback_t = ctypes.CFUNCTYPE(c__EA_rocprofiler_thread_trace_decoder_status_t, rocprofiler_thread_trace_decoder_record_type_t, ctypes.POINTER(None), ctypes.c_uint64, ctypes.POINTER(None))
rocprof_trace_decoder_isa_callback_t = ctypes.CFUNCTYPE(c__EA_rocprofiler_thread_trace_decoder_status_t, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64), struct_rocprofiler_thread_trace_decoder_pc_t, ctypes.POINTER(None))
rocprof_trace_decoder_se_data_callback_t = ctypes.CFUNCTYPE(ctypes.c_uint64, ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte)), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(None))
try:
rocprof_trace_decoder_parse_data = _libraries['FIXME_STUB'].rocprof_trace_decoder_parse_data
rocprof_trace_decoder_parse_data.restype = rocprofiler_thread_trace_decoder_status_t
rocprof_trace_decoder_parse_data.argtypes = [rocprof_trace_decoder_se_data_callback_t, rocprof_trace_decoder_trace_callback_t, rocprof_trace_decoder_isa_callback_t, ctypes.POINTER(None)]
except AttributeError:
pass
try:
rocprof_trace_decoder_get_info_string = _libraries['FIXME_STUB'].rocprof_trace_decoder_get_info_string
rocprof_trace_decoder_get_info_string.restype = ctypes.POINTER(ctypes.c_char)
rocprof_trace_decoder_get_info_string.argtypes = [rocprofiler_thread_trace_decoder_info_t]
except AttributeError:
pass
try:
rocprof_trace_decoder_get_status_string = _libraries['FIXME_STUB'].rocprof_trace_decoder_get_status_string
rocprof_trace_decoder_get_status_string.restype = ctypes.POINTER(ctypes.c_char)
rocprof_trace_decoder_get_status_string.argtypes = [rocprofiler_thread_trace_decoder_status_t]
except AttributeError:
pass
rocprofiler_thread_trace_decoder_debug_callback_t = ctypes.CFUNCTYPE(None, ctypes.c_int64, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), ctypes.POINTER(None))
rocprof_trace_decoder_trace_callback_t = ctypes.CFUNCTYPE(rocprofiler_thread_trace_decoder_status_t, enum_rocprofiler_thread_trace_decoder_record_type_t, ctypes.c_void_p, ctypes.c_uint64, ctypes.c_void_p)
class struct_rocprofiler_thread_trace_decoder_pc_t(Struct): pass
uint64_t = ctypes.c_uint64
try:
rocprof_trace_decoder_dump_data = _libraries['FIXME_STUB'].rocprof_trace_decoder_dump_data
rocprof_trace_decoder_dump_data.restype = rocprofiler_thread_trace_decoder_status_t
rocprof_trace_decoder_dump_data.argtypes = [ctypes.POINTER(ctypes.c_char), uint64_t, rocprofiler_thread_trace_decoder_debug_callback_t, ctypes.POINTER(None)]
except AttributeError:
pass
class union_rocprof_trace_decoder_gfx9_header_t(Union):
pass
class struct_rocprof_trace_decoder_gfx9_header_t_0(Structure):
pass
struct_rocprof_trace_decoder_gfx9_header_t_0._pack_ = 1 # source:False
struct_rocprof_trace_decoder_gfx9_header_t_0._fields_ = [
('legacy_version', ctypes.c_uint64, 13),
('gfx9_version2', ctypes.c_uint64, 3),
('DSIMDM', ctypes.c_uint64, 4),
('DCU', ctypes.c_uint64, 5),
('reserved1', ctypes.c_uint64, 1),
('SEID', ctypes.c_uint64, 6),
('reserved2', ctypes.c_uint64, 32),
struct_rocprofiler_thread_trace_decoder_pc_t._fields_ = [
('address', uint64_t),
('code_object_id', uint64_t),
]
rocprof_trace_decoder_isa_callback_t = ctypes.CFUNCTYPE(rocprofiler_thread_trace_decoder_status_t, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64), struct_rocprofiler_thread_trace_decoder_pc_t, ctypes.c_void_p)
rocprof_trace_decoder_se_data_callback_t = ctypes.CFUNCTYPE(ctypes.c_uint64, ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte)), ctypes.POINTER(ctypes.c_uint64), ctypes.c_void_p)
# rocprofiler_thread_trace_decoder_status_t rocprof_trace_decoder_parse_data(rocprof_trace_decoder_se_data_callback_t se_data_callback, rocprof_trace_decoder_trace_callback_t trace_callback, rocprof_trace_decoder_isa_callback_t isa_callback, void *userdata)
try: (rocprof_trace_decoder_parse_data:=dll.rocprof_trace_decoder_parse_data).restype, rocprof_trace_decoder_parse_data.argtypes = rocprofiler_thread_trace_decoder_status_t, [rocprof_trace_decoder_se_data_callback_t, rocprof_trace_decoder_trace_callback_t, rocprof_trace_decoder_isa_callback_t, ctypes.c_void_p]
except AttributeError: pass
union_rocprof_trace_decoder_gfx9_header_t._pack_ = 1 # source:False
union_rocprof_trace_decoder_gfx9_header_t._anonymous_ = ('_0',)
enum_rocprofiler_thread_trace_decoder_info_t = CEnum(ctypes.c_uint32)
ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE = enum_rocprofiler_thread_trace_decoder_info_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE', 0)
ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST = enum_rocprofiler_thread_trace_decoder_info_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST', 1)
ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE = enum_rocprofiler_thread_trace_decoder_info_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE', 2)
ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE = enum_rocprofiler_thread_trace_decoder_info_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE', 3)
ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST = enum_rocprofiler_thread_trace_decoder_info_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST', 4)
rocprofiler_thread_trace_decoder_info_t = enum_rocprofiler_thread_trace_decoder_info_t
# const char *rocprof_trace_decoder_get_info_string(rocprofiler_thread_trace_decoder_info_t info)
try: (rocprof_trace_decoder_get_info_string:=dll.rocprof_trace_decoder_get_info_string).restype, rocprof_trace_decoder_get_info_string.argtypes = ctypes.POINTER(ctypes.c_char), [rocprofiler_thread_trace_decoder_info_t]
except AttributeError: pass
# const char *rocprof_trace_decoder_get_status_string(rocprofiler_thread_trace_decoder_status_t status)
try: (rocprof_trace_decoder_get_status_string:=dll.rocprof_trace_decoder_get_status_string).restype, rocprof_trace_decoder_get_status_string.argtypes = ctypes.POINTER(ctypes.c_char), [rocprofiler_thread_trace_decoder_status_t]
except AttributeError: pass
rocprofiler_thread_trace_decoder_debug_callback_t = ctypes.CFUNCTYPE(None, ctypes.c_int64, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), ctypes.c_void_p)
# rocprofiler_thread_trace_decoder_status_t rocprof_trace_decoder_dump_data(const char *data, uint64_t data_size, rocprofiler_thread_trace_decoder_debug_callback_t cb, void *userdata)
try: (rocprof_trace_decoder_dump_data:=dll.rocprof_trace_decoder_dump_data).restype, rocprof_trace_decoder_dump_data.argtypes = rocprofiler_thread_trace_decoder_status_t, [ctypes.POINTER(ctypes.c_char), uint64_t, rocprofiler_thread_trace_decoder_debug_callback_t, ctypes.c_void_p]
except AttributeError: pass
class union_rocprof_trace_decoder_gfx9_header_t(ctypes.Union): pass
class union_rocprof_trace_decoder_gfx9_header_t_0(Struct): pass
union_rocprof_trace_decoder_gfx9_header_t_0._fields_ = [
('legacy_version', uint64_t,13),
('gfx9_version2', uint64_t,3),
('DSIMDM', uint64_t,4),
('DCU', uint64_t,5),
('reserved1', uint64_t,1),
('SEID', uint64_t,6),
('reserved2', uint64_t,32),
]
union_rocprof_trace_decoder_gfx9_header_t._anonymous_ = ['_0']
union_rocprof_trace_decoder_gfx9_header_t._fields_ = [
('_0', struct_rocprof_trace_decoder_gfx9_header_t_0),
('raw', ctypes.c_uint64),
('_0', union_rocprof_trace_decoder_gfx9_header_t_0),
('raw', uint64_t),
]
rocprof_trace_decoder_gfx9_header_t = union_rocprof_trace_decoder_gfx9_header_t
class union_rocprof_trace_decoder_instrument_enable_t(Union):
pass
class struct_rocprof_trace_decoder_instrument_enable_t_0(Structure):
pass
struct_rocprof_trace_decoder_instrument_enable_t_0._pack_ = 1 # source:False
struct_rocprof_trace_decoder_instrument_enable_t_0._fields_ = [
('char1', ctypes.c_uint32, 8),
('char2', ctypes.c_uint32, 8),
('char3', ctypes.c_uint32, 8),
('char4', ctypes.c_uint32, 8),
class union_rocprof_trace_decoder_instrument_enable_t(ctypes.Union): pass
class union_rocprof_trace_decoder_instrument_enable_t_0(Struct): pass
union_rocprof_trace_decoder_instrument_enable_t_0._fields_ = [
('char1', ctypes.c_uint32,8),
('char2', ctypes.c_uint32,8),
('char3', ctypes.c_uint32,8),
('char4', ctypes.c_uint32,8),
]
union_rocprof_trace_decoder_instrument_enable_t._pack_ = 1 # source:False
union_rocprof_trace_decoder_instrument_enable_t._anonymous_ = ('_0',)
union_rocprof_trace_decoder_instrument_enable_t._anonymous_ = ['_0']
union_rocprof_trace_decoder_instrument_enable_t._fields_ = [
('_0', struct_rocprof_trace_decoder_instrument_enable_t_0),
('u32All', ctypes.c_uint32),
('_0', union_rocprof_trace_decoder_instrument_enable_t_0),
('u32All', ctypes.c_uint32),
]
rocprof_trace_decoder_instrument_enable_t = union_rocprof_trace_decoder_instrument_enable_t
class union_rocprof_trace_decoder_packet_header_t(Union):
pass
class struct_rocprof_trace_decoder_packet_header_t_0(Structure):
pass
struct_rocprof_trace_decoder_packet_header_t_0._pack_ = 1 # source:False
struct_rocprof_trace_decoder_packet_header_t_0._fields_ = [
('opcode', ctypes.c_uint32, 8),
('type', ctypes.c_uint32, 4),
('data20', ctypes.c_uint32, 20),
class union_rocprof_trace_decoder_packet_header_t(ctypes.Union): pass
class union_rocprof_trace_decoder_packet_header_t_0(Struct): pass
union_rocprof_trace_decoder_packet_header_t_0._fields_ = [
('opcode', ctypes.c_uint32,8),
('type', ctypes.c_uint32,4),
('data20', ctypes.c_uint32,20),
]
union_rocprof_trace_decoder_packet_header_t._pack_ = 1 # source:False
union_rocprof_trace_decoder_packet_header_t._anonymous_ = ('_0',)
union_rocprof_trace_decoder_packet_header_t._anonymous_ = ['_0']
union_rocprof_trace_decoder_packet_header_t._fields_ = [
('_0', struct_rocprof_trace_decoder_packet_header_t_0),
('u32All', ctypes.c_uint32),
('_0', union_rocprof_trace_decoder_packet_header_t_0),
('u32All', ctypes.c_uint32),
]
rocprof_trace_decoder_packet_header_t = union_rocprof_trace_decoder_packet_header_t
enum_rocprof_trace_decoder_packet_opcode_t = CEnum(ctypes.c_uint32)
ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ = enum_rocprof_trace_decoder_packet_opcode_t.define('ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ', 4)
ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP = enum_rocprof_trace_decoder_packet_opcode_t.define('ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP', 5)
ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO = enum_rocprof_trace_decoder_packet_opcode_t.define('ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO', 6)
# values for enumeration 'rocprof_trace_decoder_packet_opcode_t'
rocprof_trace_decoder_packet_opcode_t__enumvalues = {
4: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ',
5: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP',
6: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO',
}
ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ = 4
ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP = 5
ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO = 6
rocprof_trace_decoder_packet_opcode_t = ctypes.c_uint32 # enum
rocprof_trace_decoder_packet_opcode_t = enum_rocprof_trace_decoder_packet_opcode_t
enum_rocprof_trace_decoder_agent_info_type_t = CEnum(ctypes.c_uint32)
ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ = enum_rocprof_trace_decoder_agent_info_type_t.define('ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ', 0)
ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL = enum_rocprof_trace_decoder_agent_info_type_t.define('ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL', 1)
ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST = enum_rocprof_trace_decoder_agent_info_type_t.define('ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST', 2)
# values for enumeration 'rocprof_trace_decoder_agent_info_type_t'
rocprof_trace_decoder_agent_info_type_t__enumvalues = {
0: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ',
1: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL',
2: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST',
}
ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ = 0
ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL = 1
ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST = 2
rocprof_trace_decoder_agent_info_type_t = ctypes.c_uint32 # enum
class union_rocprof_trace_decoder_codeobj_marker_tail_t(Union):
pass
class struct_rocprof_trace_decoder_codeobj_marker_tail_t_0(Structure):
pass
struct_rocprof_trace_decoder_codeobj_marker_tail_t_0._pack_ = 1 # source:False
struct_rocprof_trace_decoder_codeobj_marker_tail_t_0._fields_ = [
('isUnload', ctypes.c_uint32, 1),
('bFromStart', ctypes.c_uint32, 1),
('legacy_id', ctypes.c_uint32, 30),
rocprof_trace_decoder_agent_info_type_t = enum_rocprof_trace_decoder_agent_info_type_t
class union_rocprof_trace_decoder_codeobj_marker_tail_t(ctypes.Union): pass
class union_rocprof_trace_decoder_codeobj_marker_tail_t_0(Struct): pass
uint32_t = ctypes.c_uint32
union_rocprof_trace_decoder_codeobj_marker_tail_t_0._fields_ = [
('isUnload', uint32_t,1),
('bFromStart', uint32_t,1),
('legacy_id', uint32_t,30),
]
union_rocprof_trace_decoder_codeobj_marker_tail_t._pack_ = 1 # source:False
union_rocprof_trace_decoder_codeobj_marker_tail_t._anonymous_ = ('_0',)
union_rocprof_trace_decoder_codeobj_marker_tail_t._anonymous_ = ['_0']
union_rocprof_trace_decoder_codeobj_marker_tail_t._fields_ = [
('_0', struct_rocprof_trace_decoder_codeobj_marker_tail_t_0),
('raw', ctypes.c_uint32),
('_0', union_rocprof_trace_decoder_codeobj_marker_tail_t_0),
('raw', uint32_t),
]
rocprof_trace_decoder_codeobj_marker_tail_t = union_rocprof_trace_decoder_codeobj_marker_tail_t
enum_rocprof_trace_decoder_codeobj_marker_type_t = CEnum(ctypes.c_uint32)
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL = enum_rocprof_trace_decoder_codeobj_marker_type_t.define('ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL', 0)
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO = enum_rocprof_trace_decoder_codeobj_marker_type_t.define('ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO', 1)
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO = enum_rocprof_trace_decoder_codeobj_marker_type_t.define('ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO', 2)
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI = enum_rocprof_trace_decoder_codeobj_marker_type_t.define('ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI', 3)
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI = enum_rocprof_trace_decoder_codeobj_marker_type_t.define('ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI', 4)
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO = enum_rocprof_trace_decoder_codeobj_marker_type_t.define('ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO', 5)
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI = enum_rocprof_trace_decoder_codeobj_marker_type_t.define('ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI', 6)
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST = enum_rocprof_trace_decoder_codeobj_marker_type_t.define('ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST', 7)
# values for enumeration 'rocprof_trace_decoder_codeobj_marker_type_t'
rocprof_trace_decoder_codeobj_marker_type_t__enumvalues = {
0: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL',
1: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO',
2: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO',
3: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI',
4: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI',
5: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO',
6: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI',
7: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST',
}
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL = 0
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO = 1
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO = 2
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI = 3
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI = 4
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO = 5
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI = 6
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST = 7
rocprof_trace_decoder_codeobj_marker_type_t = ctypes.c_uint32 # enum
__all__ = \
['ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST',
'ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST',
'ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE',
'ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE',
'ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE',
'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM',
'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT',
'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL',
'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST',
'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL',
'ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO',
'ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ',
'ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP',
'c__EA_rocprofiler_thread_trace_decoder_status_t',
'rocprof_trace_decoder_agent_info_type_t',
'rocprof_trace_decoder_codeobj_marker_tail_t',
'rocprof_trace_decoder_codeobj_marker_type_t',
'rocprof_trace_decoder_dump_data',
'rocprof_trace_decoder_get_info_string',
'rocprof_trace_decoder_get_status_string',
'rocprof_trace_decoder_gfx9_header_t',
'rocprof_trace_decoder_instrument_enable_t',
'rocprof_trace_decoder_isa_callback_t',
'rocprof_trace_decoder_packet_header_t',
'rocprof_trace_decoder_packet_opcode_t',
'rocprof_trace_decoder_parse_data',
'rocprof_trace_decoder_se_data_callback_t',
'rocprof_trace_decoder_trace_callback_t',
'rocprofiler_thread_trace_decoder_debug_callback_t',
'rocprofiler_thread_trace_decoder_info_t',
'rocprofiler_thread_trace_decoder_inst_category_t',
'rocprofiler_thread_trace_decoder_inst_t',
'rocprofiler_thread_trace_decoder_occupancy_t',
'rocprofiler_thread_trace_decoder_pc_t',
'rocprofiler_thread_trace_decoder_perfevent_t',
'rocprofiler_thread_trace_decoder_realtime_t',
'rocprofiler_thread_trace_decoder_record_type_t',
'rocprofiler_thread_trace_decoder_shaderdata_flags_t',
'rocprofiler_thread_trace_decoder_shaderdata_t',
'rocprofiler_thread_trace_decoder_status_t',
'rocprofiler_thread_trace_decoder_status_t__enumvalues',
'rocprofiler_thread_trace_decoder_wave_state_t',
'rocprofiler_thread_trace_decoder_wave_t',
'rocprofiler_thread_trace_decoder_wstate_type_t',
'struct_rocprof_trace_decoder_codeobj_marker_tail_t_0',
'struct_rocprof_trace_decoder_gfx9_header_t_0',
'struct_rocprof_trace_decoder_instrument_enable_t_0',
'struct_rocprof_trace_decoder_packet_header_t_0',
'struct_rocprofiler_thread_trace_decoder_inst_t',
'struct_rocprofiler_thread_trace_decoder_occupancy_t',
'struct_rocprofiler_thread_trace_decoder_pc_t',
'struct_rocprofiler_thread_trace_decoder_perfevent_t',
'struct_rocprofiler_thread_trace_decoder_realtime_t',
'struct_rocprofiler_thread_trace_decoder_shaderdata_t',
'struct_rocprofiler_thread_trace_decoder_wave_state_t',
'struct_rocprofiler_thread_trace_decoder_wave_t', 'uint64_t',
'union_rocprof_trace_decoder_codeobj_marker_tail_t',
'union_rocprof_trace_decoder_gfx9_header_t',
'union_rocprof_trace_decoder_instrument_enable_t',
'union_rocprof_trace_decoder_packet_header_t']
rocprof_trace_decoder_codeobj_marker_type_t = enum_rocprof_trace_decoder_codeobj_marker_type_t
rocprofiler_thread_trace_decoder_pc_t = struct_rocprofiler_thread_trace_decoder_pc_t
class struct_rocprofiler_thread_trace_decoder_perfevent_t(Struct): pass
int64_t = ctypes.c_int64
uint16_t = ctypes.c_uint16
uint8_t = ctypes.c_ubyte
struct_rocprofiler_thread_trace_decoder_perfevent_t._fields_ = [
('time', int64_t),
('events0', uint16_t),
('events1', uint16_t),
('events2', uint16_t),
('events3', uint16_t),
('CU', uint8_t),
('bank', uint8_t),
]
rocprofiler_thread_trace_decoder_perfevent_t = struct_rocprofiler_thread_trace_decoder_perfevent_t
class struct_rocprofiler_thread_trace_decoder_occupancy_t(Struct): pass
struct_rocprofiler_thread_trace_decoder_occupancy_t._fields_ = [
('pc', rocprofiler_thread_trace_decoder_pc_t),
('time', uint64_t),
('reserved', uint8_t),
('cu', uint8_t),
('simd', uint8_t),
('wave_id', uint8_t),
('start', uint32_t,1),
('_rsvd', uint32_t,31),
]
rocprofiler_thread_trace_decoder_occupancy_t = struct_rocprofiler_thread_trace_decoder_occupancy_t
enum_rocprofiler_thread_trace_decoder_wstate_type_t = CEnum(ctypes.c_uint32)
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY = enum_rocprofiler_thread_trace_decoder_wstate_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY', 0)
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE = enum_rocprofiler_thread_trace_decoder_wstate_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE', 1)
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC = enum_rocprofiler_thread_trace_decoder_wstate_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC', 2)
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT = enum_rocprofiler_thread_trace_decoder_wstate_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT', 3)
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL = enum_rocprofiler_thread_trace_decoder_wstate_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL', 4)
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST = enum_rocprofiler_thread_trace_decoder_wstate_type_t.define('ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST', 5)
rocprofiler_thread_trace_decoder_wstate_type_t = enum_rocprofiler_thread_trace_decoder_wstate_type_t
class struct_rocprofiler_thread_trace_decoder_wave_state_t(Struct): pass
int32_t = ctypes.c_int32
struct_rocprofiler_thread_trace_decoder_wave_state_t._fields_ = [
('type', int32_t),
('duration', int32_t),
]
rocprofiler_thread_trace_decoder_wave_state_t = struct_rocprofiler_thread_trace_decoder_wave_state_t
enum_rocprofiler_thread_trace_decoder_inst_category_t = CEnum(ctypes.c_uint32)
ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE', 0)
ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM', 1)
ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU', 2)
ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM', 3)
ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT', 4)
ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS', 5)
ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU', 6)
ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP', 7)
ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT', 8)
ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED', 9)
ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT', 10)
ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE', 11)
ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH', 12)
ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST = enum_rocprofiler_thread_trace_decoder_inst_category_t.define('ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST', 13)
rocprofiler_thread_trace_decoder_inst_category_t = enum_rocprofiler_thread_trace_decoder_inst_category_t
class struct_rocprofiler_thread_trace_decoder_inst_t(Struct): pass
struct_rocprofiler_thread_trace_decoder_inst_t._fields_ = [
('category', uint32_t,8),
('stall', uint32_t,24),
('duration', int32_t),
('time', int64_t),
('pc', rocprofiler_thread_trace_decoder_pc_t),
]
rocprofiler_thread_trace_decoder_inst_t = struct_rocprofiler_thread_trace_decoder_inst_t
class struct_rocprofiler_thread_trace_decoder_wave_t(Struct): pass
struct_rocprofiler_thread_trace_decoder_wave_t._fields_ = [
('cu', uint8_t),
('simd', uint8_t),
('wave_id', uint8_t),
('contexts', uint8_t),
('_rsvd1', uint32_t),
('_rsvd2', uint32_t),
('_rsvd3', uint32_t),
('begin_time', int64_t),
('end_time', int64_t),
('timeline_size', uint64_t),
('instructions_size', uint64_t),
('timeline_array', ctypes.POINTER(rocprofiler_thread_trace_decoder_wave_state_t)),
('instructions_array', ctypes.POINTER(rocprofiler_thread_trace_decoder_inst_t)),
]
rocprofiler_thread_trace_decoder_wave_t = struct_rocprofiler_thread_trace_decoder_wave_t
class struct_rocprofiler_thread_trace_decoder_realtime_t(Struct): pass
struct_rocprofiler_thread_trace_decoder_realtime_t._fields_ = [
('shader_clock', int64_t),
('realtime_clock', uint64_t),
('reserved', uint64_t),
]
rocprofiler_thread_trace_decoder_realtime_t = struct_rocprofiler_thread_trace_decoder_realtime_t
enum_rocprofiler_thread_trace_decoder_shaderdata_flags_t = CEnum(ctypes.c_uint32)
ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM = enum_rocprofiler_thread_trace_decoder_shaderdata_flags_t.define('ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM', 0)
ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV = enum_rocprofiler_thread_trace_decoder_shaderdata_flags_t.define('ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV', 1)
rocprofiler_thread_trace_decoder_shaderdata_flags_t = enum_rocprofiler_thread_trace_decoder_shaderdata_flags_t
class struct_rocprofiler_thread_trace_decoder_shaderdata_t(Struct): pass
struct_rocprofiler_thread_trace_decoder_shaderdata_t._fields_ = [
('time', int64_t),
('value', uint64_t),
('cu', uint8_t),
('simd', uint8_t),
('wave_id', uint8_t),
('flags', uint8_t),
('reserved', uint32_t),
]
rocprofiler_thread_trace_decoder_shaderdata_t = struct_rocprofiler_thread_trace_decoder_shaderdata_t
rocprofiler_thread_trace_decoder_record_type_t = enum_rocprofiler_thread_trace_decoder_record_type_t
tinygrad/runtime/autogen/sqtt.py
+720 -1644
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/vfio.py
+503 -812
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/autogen/webgpu.py
+3570 -6658
View File
File diff suppressed because it is too large Load Diff
tinygrad/runtime/ops_amd.py
+8 -12
View File
@@ -188,14 +188,9 @@ class AMDComputeQueue(HWQueue):
### SQTT ###
def sqtt_setup_exec(self, prg, global_size):
self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_pipeline_bind(
_0=sqtt.union_rgp_sqtt_marker_pipeline_bind_0(_0=sqtt.struct_rgp_sqtt_marker_pipeline_bind_0_0(
identifier=sqtt.RGP_SQTT_MARKER_IDENTIFIER_BIND_PIPELINE, bind_point=(__BIND_POINT_COMPUTE:=1))),
_1=sqtt.union_rgp_sqtt_marker_pipeline_bind_1(api_pso_hash=data64_le(prg.libhash[0]))))
self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_event(
_0=sqtt.union_rgp_sqtt_marker_event_0(_0=sqtt.struct_rgp_sqtt_marker_event_0_0(has_thread_dims=1)),
_2=sqtt.union_rgp_sqtt_marker_event_2(cmd_id=next(prg.dev.sqtt_next_cmd_id))), *global_size)
self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_pipeline_bind(identifier=sqtt.RGP_SQTT_MARKER_IDENTIFIER_BIND_PIPELINE,
bind_point=(__BIND_POINT_COMPUTE:=1), api_pso_hash=data64_le(prg.libhash[0])))
self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_event(has_thread_dims=1, cmd_id=next(prg.dev.sqtt_next_cmd_id)), *global_size)
se_cap = max(prod([x if isinstance(x, int) else 1 for x in global_size]) // 4, 1) // 32
for xcc in range(self.dev.xccs):
@@ -594,7 +589,8 @@ class AMDProgram(HCQProgram):
self.dev.synchronize()
for se, buf in enumerate(self.dev.sqtt_buffers):
wptr = ((self.dev.sqtt_wptrs.cpu_view().view(fmt='I')[se]&0x1FFFFFFF)-(((buf.va_addr//32)&0x1FFFFFFF) if self.dev.target[0] == 11 else 0))*32
wptr = (self.dev.sqtt_wptrs.cpu_view().view(fmt='I')[se] & 0x1FFFFFFF) * 32
if self.dev.target[:2] == (11, 0): wptr -= ((buf.va_addr // 32) & 0x1FFFFFFF) * 32
if DEBUG >= 5: print(f'\t{self.dev.device}: SE {se} blob size {wptr:#x}')
assert wptr >= 0 and wptr <= buf.size, f"{wptr} > {buf.size}, should never happen"
@@ -785,7 +781,7 @@ class PCIIface(PCIIfaceBase):
gpus:ClassVar[list[str]] = []
def __init__(self, dev, dev_id):
super().__init__(dev, dev_id, vendor=0x1002, devices=[0x744c, 0x7480, 0x7550, 0x7590], bars=[0, 2, 5], vram_bar=0,
super().__init__(dev, dev_id, vendor=0x1002, devices=[(0xffff, [0x744c, 0x7480, 0x7550, 0x7590])], bars=[0, 2, 5], vram_bar=0,
va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size)
self._setup_adev(self.pci_dev)
self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
@@ -911,8 +907,8 @@ class AMDDevice(HCQCompiled):
max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20))
compilers:list[CompilerPairT] = [(functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch)),
(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch))]
compilers:list[CompilerPairT] = [(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch)),
(functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch))]
super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),
tinygrad/runtime/ops_cuda.py
+2 -2
View File
@@ -11,7 +11,7 @@ if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore
def check(status):
if status != 0:
error = ctypes.string_at(init_c_var(ctypes.POINTER(ctypes.c_char)(), lambda x: cuda.cuGetErrorString(status, ctypes.byref(x)))).decode()
error = ctypes.string_at(init_c_var(ctypes.c_char_p(), lambda x: cuda.cuGetErrorString(status, x))).decode()
raise RuntimeError(f"CUDA Error {status}, {error}")
def encode_args(args, vals) -> tuple[ctypes.Structure, ctypes.Array]:
@@ -42,7 +42,7 @@ class CUDAProgram:
status = cuda.cuModuleLoadData(ctypes.byref(self.module), lib)
if status != 0:
del self.module
raise RuntimeError(f"module load failed with status code {status}: {cuda.cudaError_enum__enumvalues[status]}")
raise RuntimeError(f"module load failed with status code {status}: {cuda.CUresult.get(status)}")
check(cuda.cuModuleGetFunction(ctypes.byref(prg := cuda.CUfunction()), self.module, name.encode("utf-8")))
self.prg = prg
if self.smem > 0: check(cuda.cuFuncSetAttribute(self.prg, cuda.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, self.smem))
tinygrad/runtime/ops_nv.py
+37 -32
View File
@@ -12,33 +12,24 @@ from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.cstyle import NVRenderer
from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, NVPTXCompiler, NVCompiler
from tinygrad.runtime.support.compiler_mesa import NAKCompiler
from tinygrad.runtime.autogen import nv_gpu, pci, mesa
from tinygrad.runtime.autogen import nv_570, nv_580, pci, mesa
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.support.nv.nvdev import NVDev, NVMemoryManager
from tinygrad.runtime.support.system import System, PCIIfaceBase, MAP_FIXED
from tinygrad.renderer.nir import NAKRenderer
if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl # noqa: F401 # pylint: disable=unused-import
nv_gpu = nv_570 # default to 570
def get_error_str(status): return f"{status}: {nv_gpu.nv_status_codes.get(status, 'Unknown error')}"
NV_PFAULT_FAULT_TYPE = {dt:name for name,dt in nv_gpu.__dict__.items() if name.startswith("NV_PFAULT_FAULT_TYPE_")}
NV_PFAULT_ACCESS_TYPE = {dt:name.split("_")[-1] for name,dt in nv_gpu.__dict__.items() if name.startswith("NV_PFAULT_ACCESS_TYPE_")}
def nv_iowr(fd:FileIOInterface, nr, args):
ret = fd.ioctl((3 << 30) | (ctypes.sizeof(args) & 0x1FFF) << 16 | (ord('F') & 0xFF) << 8 | (nr & 0xFF), args)
def nv_iowr(fd:FileIOInterface, nr, args, cmd=None):
ret = fd.ioctl(cmd or ((3 << 30) | (ctypes.sizeof(args) & 0x1FFF) << 16 | (ord('F') & 0xFF) << 8 | (nr & 0xFF)), args)
if ret != 0: raise RuntimeError(f"ioctl returned {ret}")
def uvm_ioctl(cmd, sttyp, fd:FileIOInterface, **kwargs):
ret = fd.ioctl(cmd, made:=sttyp(**kwargs))
if ret != 0: raise RuntimeError(f"ioctl(uvm) returned {ret}")
if made.rmStatus != 0: raise RuntimeError(f"uvm_ioctl returned {get_error_str(made.rmStatus)}")
return made
def make_uvm_type():
return type("NVUVM", (object,), {name.replace("UVM_", "").lower(): functools.partial(uvm_ioctl, dt, getattr(nv_gpu, name+"_PARAMS"))
for name,dt in nv_gpu.__dict__.items() if name.startswith("UVM_") and nv_gpu.__dict__.get(name+"_PARAMS")})
uvm = make_uvm_type()
class QMD:
fields: dict[str, dict[str, tuple[int, int]]] = {}
@@ -313,12 +304,20 @@ class NVKIface:
def __init__(self, dev, device_id):
if NVKIface.root is None:
global nv_gpu
NVKIface.fd_ctl = FileIOInterface("/dev/nvidiactl", os.O_RDWR | os.O_CLOEXEC)
NVKIface.fd_uvm = FileIOInterface("/dev/nvidia-uvm", os.O_RDWR | os.O_CLOEXEC)
self.fd_uvm_2 = FileIOInterface("/dev/nvidia-uvm", os.O_RDWR | os.O_CLOEXEC)
NVKIface.root = self.rm_alloc(0, nv_gpu.NV01_ROOT_CLIENT, None, root=0)
uvm.initialize(self.fd_uvm)
with contextlib.suppress(RuntimeError): uvm.mm_initialize(self.fd_uvm_2, uvmFd=self.fd_uvm.fd) # this error is okay, CUDA hits it too
drvver = self.rm_control(self.root, nv_gpu.NV0000_CTRL_CMD_SYSTEM_GET_BUILD_VERSION_V2, nv_gpu.NV0000_CTRL_SYSTEM_GET_BUILD_VERSION_V2_PARAMS())
if int(drvver.driverVersionBuffer.decode().split('.')[0], 10) >= 580: nv_gpu = nv_580
self.uvm(nv_gpu.UVM_INITIALIZE, nv_gpu.UVM_INITIALIZE_PARAMS())
# this error is okay, CUDA hits it too
with contextlib.suppress(RuntimeError): self.uvm(nv_gpu.UVM_MM_INITIALIZE, nv_gpu.UVM_MM_INITIALIZE_PARAMS(uvmFd=self.fd_uvm.fd), self.fd_uvm_2)
nv_iowr(NVKIface.fd_ctl, nv_gpu.NV_ESC_CARD_INFO, gpus_info:=(nv_gpu.nv_ioctl_card_info_t*64)())
NVKIface.gpus_info = hcq_filter_visible_devices(gpus_info)
@@ -346,9 +345,14 @@ class NVKIface:
if made.status != 0: raise RuntimeError(f"rm_control returned {get_error_str(made.status)}")
return params
def uvm(self, cmd, params, fd=None):
nv_iowr(fd or self.fd_uvm, None, params, cmd=cmd)
if params.rmStatus != 0: raise RuntimeError(f"uvm returned {get_error_str(params.rmStatus)}")
def setup_usermode(self):
clsinfo = self.rm_control(self.dev.nvdevice, nv_gpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST, nv_gpu.NV0080_CTRL_GPU_GET_CLASSLIST_PARAMS(numClasses=100,
classList=mv_address(classlist:=memoryview(bytearray(100 * 4)).cast('I'))))
clsnum = self.rm_control(self.dev.nvdevice, nv_gpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST, nv_gpu.NV0080_CTRL_GPU_GET_CLASSLIST_PARAMS(numClasses=0))
clsinfo = self.rm_control(self.dev.nvdevice, nv_gpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST, nv_gpu.NV0080_CTRL_GPU_GET_CLASSLIST_PARAMS(
numClasses=clsnum.numClasses, classList=mv_address(classlist:=memoryview(bytearray(clsnum.numClasses * 4)).cast('I'))))
self.nvclasses = {classlist[i] for i in range(clsinfo.numClasses)}
self.usermode_class:int = next(c for c in [nv_gpu.HOPPER_USERMODE_A, nv_gpu.TURING_USERMODE_A] if c in self.nvclasses)
self.gpfifo_class:int = next(c for c in [nv_gpu.BLACKWELL_CHANNEL_GPFIFO_A, nv_gpu.AMPERE_CHANNEL_GPFIFO_A] if c in self.nvclasses)
@@ -363,16 +367,17 @@ class NVKIface:
flags=nv_gpu.NV2080_GPU_CMD_GPU_GET_GID_FLAGS_FORMAT_BINARY, length=16))
self.gpu_uuid = nv_gpu.struct_nv_uuid(uuid=(ctypes.c_ubyte*16)(*[raw_uuid.data[i] for i in range(16)]))
uvm.register_gpu(self.fd_uvm, rmCtrlFd=-1, gpu_uuid=self.gpu_uuid)
uvm.register_gpu_vaspace(self.fd_uvm, gpuUuid=self.gpu_uuid, rmCtrlFd=self.fd_ctl.fd, hClient=self.root, hVaSpace=vaspace)
self.uvm(nv_gpu.UVM_REGISTER_GPU, nv_gpu.UVM_REGISTER_GPU_PARAMS(rmCtrlFd=-1, gpu_uuid=self.gpu_uuid))
self.uvm(nv_gpu.UVM_REGISTER_GPU_VASPACE, nv_gpu.UVM_REGISTER_GPU_VASPACE_PARAMS(
gpuUuid=self.gpu_uuid, rmCtrlFd=self.fd_ctl.fd, hClient=self.root, hVaSpace=vaspace))
for dev in cast(list[NVDevice], [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, NVDevice) and not d.is_nvd()]):
try: uvm.enable_peer_access(self.fd_uvm, gpuUuidA=self.gpu_uuid, gpuUuidB=dev.iface.gpu_uuid)
try: self.uvm(nv_gpu.UVM_ENABLE_PEER_ACCESS, nv_gpu.UVM_ENABLE_PEER_ACCESS_PARAMS(gpuUuidA=self.gpu_uuid, gpuUuidB=dev.iface.gpu_uuid))
except RuntimeError as e: raise RuntimeError(f"{e}. Make sure GPUs #{self.gpu_minor} & #{dev.iface.gpu_minor} have P2P enabled.") from e
def setup_gpfifo_vm(self, gpfifo):
uvm.register_channel(self.fd_uvm, gpuUuid=self.gpu_uuid, rmCtrlFd=self.fd_ctl.fd, hClient=self.root,
hChannel=gpfifo, base=self._alloc_gpu_vaddr(0x4000000, force_low=True), length=0x4000000)
self.uvm(nv_gpu.UVM_REGISTER_CHANNEL, nv_gpu.UVM_REGISTER_CHANNEL_PARAMS(gpuUuid=self.gpu_uuid, rmCtrlFd=self.fd_ctl.fd, hClient=self.root,
hChannel=gpfifo, base=self._alloc_gpu_vaddr(0x4000000, force_low=True), length=0x4000000))
def _new_gpu_fd(self):
fd_dev = FileIOInterface(f"/dev/nvidia{NVKIface.gpus_info[self.device_id].minor_number}", os.O_RDWR | os.O_CLOEXEC)
@@ -431,17 +436,17 @@ class NVKIface:
nv_iowr(self.fd_ctl, nv_gpu.NV_ESC_RM_FREE, made)
if made.status != 0: raise RuntimeError(f"_gpu_free returned {get_error_str(made.status)}")
uvm.free(self.fd_uvm, base=cast(int, mem.va_addr), length=mem.size)
self.uvm(nv_gpu.UVM_FREE, nv_gpu.UVM_FREE_PARAMS(base=cast(int, mem.va_addr), length=mem.size))
if mem.meta.has_cpu_mapping: FileIOInterface.munmap(cast(int, mem.va_addr), mem.size)
def _gpu_uvm_map(self, va_base, size, mem_handle, create_range=True, has_cpu_mapping=False) -> HCQBuffer:
if create_range: uvm.create_external_range(self.fd_uvm, base=va_base, length=size)
attrs = (nv_gpu.struct_c__SA_UvmGpuMappingAttributes*256)(nv_gpu.struct_c__SA_UvmGpuMappingAttributes(gpuUuid=self.gpu_uuid, gpuMappingType=1))
if create_range: self.uvm(nv_gpu.UVM_CREATE_EXTERNAL_RANGE, nv_gpu.UVM_CREATE_EXTERNAL_RANGE_PARAMS(base=va_base, length=size))
attrs = (nv_gpu.UvmGpuMappingAttributes*256)(nv_gpu.UvmGpuMappingAttributes(gpuUuid=self.gpu_uuid, gpuMappingType=1))
# NOTE: va_addr is set to make rawbufs compatible with HCQBuffer protocol.
return HCQBuffer(va_base, size, meta=uvm.map_external_allocation(self.fd_uvm, base=va_base, length=size, rmCtrlFd=self.fd_ctl.fd,
hClient=self.root, hMemory=mem_handle, gpuAttributesCount=1, perGpuAttributes=attrs, mapped_gpu_ids=[self.gpu_uuid],
has_cpu_mapping=has_cpu_mapping), view=MMIOInterface(va_base, size, fmt='B') if has_cpu_mapping else None, owner=self.dev)
self.uvm(nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION, uvm_map:=nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS(base=va_base, length=size,
rmCtrlFd=self.fd_ctl.fd, hClient=self.root, hMemory=mem_handle, gpuAttributesCount=1, perGpuAttributes=attrs, mapped_gpu_ids=[self.gpu_uuid],
has_cpu_mapping=has_cpu_mapping))
return HCQBuffer(va_base, size, meta=uvm_map, view=MMIOInterface(va_base, size, fmt='B') if has_cpu_mapping else None, owner=self.dev)
def map(self, mem:HCQBuffer):
if mem.owner is not None and mem.owner._is_cpu():
@@ -456,8 +461,8 @@ class PCIIface(PCIIfaceBase):
gpus:ClassVar[list[str]] = []
def __init__(self, dev, dev_id):
super().__init__(dev, dev_id, vendor=0x10de, devices=[0x2204, 0x2684, 0x2b85], bars=[0, 1], vram_bar=1,
va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
super().__init__(dev, dev_id, vendor=0x10de, devices=[(0xff00, [0x2200, 0x2400, 0x2500, 0x2600, 0x2700, 0x2800, 0x2b00, 0x2c00, 0x2d00, 0x2f00])],
bars=[0, 1], vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
if not OSX: System.reserve_hugepages(64)
self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
tinygrad/runtime/ops_qcom.py
+14 -8
View File
@@ -1,5 +1,5 @@
from __future__ import annotations
import os, ctypes, functools, mmap, struct, array, math, sys, weakref
import os, ctypes, functools, mmap, struct, array, math, sys, weakref, contextlib
assert sys.platform != 'win32'
from types import SimpleNamespace
from typing import Any, cast
@@ -9,7 +9,8 @@ from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
from tinygrad.runtime.autogen import kgsl, adreno
from tinygrad.runtime.ops_cl import CLCompiler, CLDevice
from tinygrad.renderer.cstyle import QCOMRenderer
from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile, lo32, PROFILE, colored
from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile, lo32, PROFILE
from tinygrad.runtime.support.system import System
if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl # noqa: F401 # pylint: disable=unused-import
BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2
@@ -320,13 +321,15 @@ class QCOMAllocator(HCQAllocatorBase):
self.dev.synchronize()
self.dev._gpu_free(opaque)
def flag(nm, val): return (val << getattr(kgsl, f"{nm}_SHIFT")) & getattr(kgsl, f"{nm}_MASK")
class QCOMDevice(HCQCompiled):
def __init__(self, device:str=""):
self.fd = FileIOInterface('/dev/kgsl-3d0', os.O_RDWR)
self.dummy_addr = cast(int, self._gpu_alloc(0x1000).va_addr)
flags = kgsl.KGSL_CONTEXT_PREAMBLE | kgsl.KGSL_CONTEXT_PWR_CONSTRAINT | kgsl.KGSL_CONTEXT_NO_FAULT_TOLERANCE | kgsl.KGSL_CONTEXT_NO_GMEM_ALLOC \
| kgsl.KGSL_CONTEXT_PRIORITY(getenv("QCOM_PRIORITY", 8)) | kgsl.KGSL_CONTEXT_PREEMPT_STYLE(kgsl.KGSL_CONTEXT_PREEMPT_STYLE_FINEGRAIN)
| flag("KGSL_CONTEXT_PRIORITY", getenv("QCOM_PRIORITY", 8)) | flag("KGSL_CONTEXT_PREEMPT_STYLE", kgsl.KGSL_CONTEXT_PREEMPT_STYLE_FINEGRAIN)
self.ctx = kgsl.IOCTL_KGSL_DRAWCTXT_CREATE(self.fd, flags=flags).drawctxt_id
self.cmd_buf = self._gpu_alloc(16 << 20)
@@ -348,17 +351,16 @@ class QCOMDevice(HCQCompiled):
# a7xx start with 730x or 'Cxxx', a8xx starts 'Exxx'
if self.gpu_id[:2] >= (7, 3): raise RuntimeError(f"Unsupported GPU: chip_id={info.chip_id:#x}")
if PROFILE and self.gpu_id[:2] < (7, 3) and int(FileIOInterface('/sys/class/kgsl/kgsl-3d0/idle_timer', os.O_RDONLY).read(), 0) < 4000000000:
print(colored("WARNING: gpu can go into suspend mode and reset timestamps. "
"Run 'echo \"4294947000\" | sudo tee /sys/class/kgsl/kgsl-3d0/idle_timer' to prevent idle state.", "yellow"))
if PROFILE and self.gpu_id[:2] < (7, 3):
System.write_sysfs("/sys/class/kgsl/kgsl-3d0/idle_timer", value="4000000000", msg="Failed to disable suspend mode", expected="4294967276")
compilers = [(QCOMRenderer, functools.partial(QCOMCompiler, device))]
super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal,
functools.partial(QCOMComputeQueue, self), None)
def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
flags |= kgsl.KGSL_MEMALIGN(alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP
if uncached: flags |= kgsl.KGSL_CACHEMODE(kgsl.KGSL_CACHEMODE_UNCACHED)
flags |= flag("KGSL_MEMALIGN", alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP
if uncached: flags |= flag("KGSL_CACHEMODE", kgsl.KGSL_CACHEMODE_UNCACHED)
alloc = kgsl.IOCTL_KGSL_GPUOBJ_ALLOC(self.fd, size=(bosz:=round_up(size, 1<<alignment_hint)), flags=flags, mmapsize=bosz)
va_addr = self.fd.mmap(0, bosz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, alloc.id * 0x1000)
@@ -376,3 +378,7 @@ class QCOMDevice(HCQCompiled):
self.synchronize()
self._gpu_free(self._stack)
self._stack = self._gpu_alloc(sz)
def _at_profile_finalize(self):
super()._at_profile_finalize()
with contextlib.suppress(RuntimeError): System.write_sysfs("/sys/class/kgsl/kgsl-3d0/idle_timer", "10", "Failed to reenable suspend mode")
tinygrad/runtime/ops_webgpu.py
+8 -10
View File
@@ -10,7 +10,7 @@ import os
WGPUDevPtr: TypeAlias = webgpu.WGPUDevice # type: ignore
WGPUBufPtr: TypeAlias = webgpu.WGPUBuffer # type: ignore
backend_types = {v: k for k, v in webgpu.WGPUBackendType__enumvalues.items() }
backend_types = {v: k for k, v in webgpu.enum_WGPUBackendType.items() }
instance = webgpu.wgpuCreateInstance(webgpu.WGPUInstanceDescriptor(features = webgpu.WGPUInstanceFeatures(timedWaitAnyEnable = True)))
@@ -38,7 +38,7 @@ def _run(async_fun, cb_info_type, cb_type, status_enum, res_idx:int|None, msg_id
cb_info = cb_info_type(nextInChain=None, mode=webgpu.WGPUCallbackMode_WaitAnyOnly, callback=cb_type(cb))
_wait(async_fun(*params, cb_info))
if result[0] != 1: raise RuntimeError(f"[{status_enum[result[0]] if status_enum else 'ERROR'}]{result[msg_idx] if msg_idx else ''}")
if result[0] != 1: raise RuntimeError(f"[{status_enum.get(result[0]) if status_enum else 'ERROR'}]{result[msg_idx] if msg_idx else ''}")
return result[res_idx] if res_idx else None
def copy_buffer_to_buffer(dev:WGPUDevPtr, src:WGPUBufPtr, src_offset:int, dst:WGPUBufPtr, dst_offset:int, size:int):
@@ -54,8 +54,8 @@ def read_buffer(dev:WGPUDevPtr, buf:WGPUBufPtr) -> memoryview:
tmp_buffer = webgpu.wgpuDeviceCreateBuffer(dev, webgpu.WGPUBufferDescriptor(size=size,
usage=webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_MapRead, mappedAtCreation=False))
copy_buffer_to_buffer(dev, buf, 0, tmp_buffer, 0, size)
_run(webgpu.wgpuBufferMapAsync2, webgpu.WGPUBufferMapCallbackInfo2, webgpu.WGPUBufferMapCallback2, webgpu.WGPUBufferMapAsyncStatus__enumvalues,
None, 0, tmp_buffer, webgpu.WGPUMapMode_Read, 0, size)
_run(webgpu.wgpuBufferMapAsync2, webgpu.WGPUBufferMapCallbackInfo2, webgpu.WGPUBufferMapCallback2, webgpu.WGPUBufferMapAsyncStatus, None, 0,
tmp_buffer, webgpu.WGPUMapMode_Read, 0, size)
void_ptr = ctypes.cast(webgpu.wgpuBufferGetConstMappedRange(tmp_buffer, 0, size), ctypes.c_void_p)
buf_copy = bytearray((ctypes.c_uint8 * size).from_address(void_ptr.value))
webgpu.wgpuBufferUnmap(tmp_buffer)
@@ -140,7 +140,7 @@ class WebGPUProgram:
compute_desc = webgpu.WGPUComputePipelineDescriptor(layout=pipeline_layout,
compute=webgpu.WGPUComputeState(module=self.prg, entryPoint=to_wgpu_str(self.name)))
pipeline_result = _run(webgpu.wgpuDeviceCreateComputePipelineAsync2, webgpu.WGPUCreateComputePipelineAsyncCallbackInfo2,
webgpu.WGPUCreateComputePipelineAsyncCallback2, webgpu.WGPUCreatePipelineAsyncStatus__enumvalues, 1, None, self.dev, compute_desc)
webgpu.WGPUCreateComputePipelineAsyncCallback2, webgpu.WGPUCreatePipelineAsyncStatus, 1, None, self.dev, compute_desc)
command_encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.dev, webgpu.WGPUCommandEncoderDescriptor())
comp_pass_desc = webgpu.WGPUComputePassDescriptor(nextInChain=None)
@@ -195,9 +195,7 @@ class WebGpuDevice(Compiled):
def __init__(self, device:str):
# Requesting an adapter
adapter_res = _run(webgpu.wgpuInstanceRequestAdapterF, webgpu.WGPURequestAdapterCallbackInfo, webgpu.WGPURequestAdapterCallback,
webgpu.WGPURequestAdapterStatus__enumvalues, 1, 2, instance,
webgpu.WGPURequestAdapterOptions(powerPreference=webgpu.WGPUPowerPreference_HighPerformance,
webgpu.WGPURequestAdapterStatus, 1, 2, instance, webgpu.WGPURequestAdapterOptions(powerPreference=webgpu.WGPUPowerPreference_HighPerformance,
backendType=backend_types.get(os.getenv("WEBGPU_BACKEND", ""), 0)))
# Get supported features
@@ -215,11 +213,11 @@ class WebGpuDevice(Compiled):
# Requesting a device
device_res = _run(webgpu.wgpuAdapterRequestDeviceF, webgpu.WGPURequestDeviceCallbackInfo, webgpu.WGPURequestDeviceCallback,
webgpu.WGPURequestDeviceStatus__enumvalues, 1, 2, adapter_res, dev_desc)
webgpu.WGPURequestDeviceStatus, 1, 2, adapter_res, dev_desc)
super().__init__(device, WebGpuAllocator(device_res), [(WGSLRenderer, Compiler)],
functools.partial(WebGPUProgram, (device_res, webgpu.WGPUFeatureName_TimestampQuery in supported)))
def synchronize(self):
_run(webgpu.wgpuQueueOnSubmittedWorkDone2, webgpu.WGPUQueueWorkDoneCallbackInfo2, webgpu.WGPUQueueWorkDoneCallback2,
webgpu.WGPUQueueWorkDoneStatus__enumvalues, None, None, webgpu.wgpuDeviceGetQueue(self.runtime.args[0][0]))
webgpu.WGPUQueueWorkDoneStatus, None, None, webgpu.wgpuDeviceGetQueue(self.runtime.args[0][0]))
tinygrad/runtime/support/am/amdev.py
+1 -3
View File
@@ -238,8 +238,6 @@ class AMDev(PCIDevImplBase):
ihdr = am.struct_ip_discovery_header.from_address(ctypes.addressof(self.bhdr) + self.bhdr.table_list[am.IP_DISCOVERY].offset)
assert self.bhdr.binary_signature == am.BINARY_SIGNATURE and ihdr.signature == am.DISCOVERY_TABLE_SIGNATURE, "discovery signatures mismatch"
# Mapping of HW IP to Discovery HW IP
hw_id_map = {am.__dict__[x]: int(y) for x,y in am.hw_id_map}
self.regs_offset:dict[int, dict[int, tuple]] = collections.defaultdict(dict)
self.ip_ver:dict[int, tuple[int, int, int]] = {}
@@ -251,7 +249,7 @@ class AMDev(PCIDevImplBase):
ip = am.struct_ip_v4.from_address(ip_offset)
ba = ((ctypes.c_uint64 if ihdr.base_addr_64_bit else ctypes.c_uint32) * ip.num_base_address).from_address(ip_offset + 8)
for hw_ip in range(1, am.MAX_HWIP):
if hw_ip in hw_id_map and hw_id_map[hw_ip] == ip.hw_id:
if hw_ip in am.hw_id_map and am.hw_id_map[hw_ip] == ip.hw_id:
self.regs_offset[hw_ip][ip.instance_number] = tuple(list(ba))
self.ip_ver[hw_ip] = (ip.major, ip.minor, ip.revision)
tinygrad/runtime/support/am/ip.py
+2 -2
View File
@@ -427,7 +427,7 @@ class AM_PSP(AM_IP):
self._wait_for_bootloader()
if DEBUG >= 2: print(f"am {self.adev.devfmt}: loading sos component: {am.psp_fw_type__enumvalues[fw]}")
if DEBUG >= 2: print(f"am {self.adev.devfmt}: loading sos component: {am.enum_psp_fw_type.get(fw)}")
self._prep_msg1(self.adev.fw.sos_fw[fw])
self.adev.reg(f"{self.reg_pref}_36").write(self.msg1_addr >> 20)
@@ -482,7 +482,7 @@ class AM_PSP(AM_IP):
def _load_ip_fw_cmd(self, fw_types:list[int], fw_bytes:memoryview):
self._prep_msg1(fw_bytes)
for fw_type in fw_types:
if DEBUG >= 2: print(f"am {self.adev.devfmt}: loading fw: {am.psp_gfx_fw_type__enumvalues[fw_type]}")
if DEBUG >= 2: print(f"am {self.adev.devfmt}: loading fw: {am.enum_psp_gfx_fw_type.get(fw_type)}")
cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_IP_FW)
cmd.cmd.cmd_load_ip_fw.fw_phy_addr_hi, cmd.cmd.cmd_load_ip_fw.fw_phy_addr_lo = data64(self.msg1_addr)
cmd.cmd.cmd_load_ip_fw.fw_size = len(fw_bytes)
tinygrad/runtime/support/autogen.py
+6 -3
View File
@@ -1,9 +1,9 @@
import ctypes.util, importlib.metadata, itertools, re, functools, os
from tinygrad.helpers import flatten, unwrap
assert importlib.metadata.version('clang')[:2] == "20", 'clang version 20 required, pip install "clang==20.1.0"'
from clang.cindex import Config, Index, CursorKind as CK, TranslationUnit as TU, LinkageKind as LK, TokenKind as ToK, TypeKind as TK
from clang.cindex import PrintingPolicy as PP, PrintingPolicyProperty as PPP, SourceRange
assert importlib.metadata.version('clang')[:2] == "20"
if not Config.loaded: Config.set_library_file(os.getenv("LIBCLANG_PATH", ctypes.util.find_library("clang-20")))
def fst(c): return next(c.get_children())
@@ -47,7 +47,6 @@ def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_e
return types[t.spelling][0]
case TK.RECORD:
# TODO: packed unions
# TODO: pragma pack support
# check for forward declaration
if t.spelling in types: types[t.spelling] = (nm:=types[t.spelling][0]), len(list(t.get_fields())) != 0
else:
@@ -56,11 +55,15 @@ def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_e
else: types[t.spelling] = (nm:=t.spelling.replace(' ', '_').replace('::', '_')), len(list(t.get_fields())) != 0
lines.append(f"class {nm}({'Struct' if decl.kind==CK.STRUCT_DECL else 'ctypes.Union'}): pass")
if typedef: lines.append(f"{typedef} = {nm}")
if (is_packed:=(CK.PACKED_ATTR in attrs(decl)) or ((N:=t.get_align()) != max([f.type.get_align() for f in t.get_fields()], default=N))):
if t.get_align() != 1:
print(f"WARNING: ignoring alignment={t.get_align()} on {t.spelling}")
is_packed = False
acnt = itertools.count().__next__
ll=[" ("+((fn:=f"'_{acnt()}'")+f", {tname(f.type, nm+fn[1:-1])}" if f.is_anonymous_record_decl() else f"'{f.spelling}', "+
tname(f.type, f'{nm}_{f.spelling}'))+(f',{f.get_bitfield_width()}' if f.is_bitfield() else '')+")," for f in t.get_fields()]
lines.extend(([f"{nm}._anonymous_ = ["+", ".join(f"'_{i}'" for i in range(n))+"]"] if (n:=acnt()) else [])+
([f"{nm}._packed_ = True"] * (CK.PACKED_ATTR in attrs(decl)))+([f"{nm}._fields_ = [",*ll,"]"] if ll else []))
([f"{nm}._packed_ = True"] * is_packed)+([f"{nm}._fields_ = [",*ll,"]"] if ll else []))
return nm
case TK.ENUM:
# TODO: C++ and GNU C have forward declared enums
tinygrad/runtime/support/c.py
+28 -12
View File
@@ -1,10 +1,13 @@
import ctypes, functools, sys
from typing import TYPE_CHECKING
from tinygrad.helpers import flatten
from _ctypes import _SimpleCData
def _do_ioctl(__idir, __base, __nr, __struct, __fd, **kwargs):
def _do_ioctl(__idir, __base, __nr, __struct, __fd, *args, __payload=None, **kwargs):
import tinygrad.runtime.support.hcq as hcq, fcntl
ioctl = __fd.ioctl if isinstance(__fd, hcq.FileIOInterface) else functools.partial(fcntl.ioctl, __fd)
if (rc:=ioctl((__idir<<30)|(ctypes.sizeof(out:=__struct(**kwargs))<<16)|(__base<<8)|__nr, out)): raise RuntimeError(f"ioctl returned {rc}")
if (rc:=ioctl((__idir<<30)|(ctypes.sizeof(out:=(__payload or __struct(*args, **kwargs)))<<16)|(__base<<8)|__nr, out)):
raise RuntimeError(f"ioctl returned {rc}")
return out
def _IO(base, nr): return functools.partial(_do_ioctl, 0, ord(base) if isinstance(base, str) else base, nr, None)
@@ -50,16 +53,30 @@ else:
@staticmethod
def _build(cls, fields):
o = 0
for n,t,b in [(f[0], f[1], f[2] if len(f) == 3 else 0) for f in fields]:
if b == 0: o = (o + 7) & ~7
m = (1 << (sz:=ctypes.sizeof(t)*8 if b == 0 else b)) - 1
def _s(self,v,m,s,b): self._data[:] = ((int.from_bytes(self._data,sys.byteorder)&~(m<<s))|((v&m)<<s)).to_bytes(len(self._data), sys.byteorder)
setattr(cls, n, property(functools.partial(lambda self,m,s:(int.from_bytes(self._data,sys.byteorder)>>s)&m,m=m,s=o),
functools.partial(_s,m=m,s=o,b=b)))
o += sz
offset = 0
for nm, ty, bf in [(f[0], f[1], f[2] if len(f) == 3 else 0) for f in fields]:
if bf == 0: offset = (offset + 7) & ~7
mask = (1 << (sz:=ctypes.sizeof(ty)*8 if bf == 0 else bf)) - 1
def fget(self, mask, off, ty): return ((int.from_bytes(self._data, sys.byteorder)>>off)&mask if issubclass(ty, _SimpleCData) else
ty.from_buffer(memoryview(self._data)[(st:=off//8):st+ctypes.sizeof(ty)]))
def fset(self, val, mask, off):
if val.__class__ is not int: val = int.from_bytes(val, sys.byteorder)
self._data[:] = (((int.from_bytes(self._data, sys.byteorder) & ~(mask<<off))|((val&mask)<<off))
.to_bytes(len(self._data), sys.byteorder))
setattr(cls, nm, property(functools.partial(fget, mask=mask, off=offset, ty=ty), functools.partial(fset, mask=mask, off=offset)))
offset += sz
type(ctypes.Structure).__setattr__(cls, '_fields_', [('_data', ctypes.c_ubyte * ((o + 7) // 8))])
def pget(ty, s): return getattr(ty, f'_packed_{s}_', getattr(ty, f'_{s}_', []))
def get_aty(anm, fs=fields): return next(f[1] for f in fs if f[0] == anm)
def get_fnms(ty): return [f[0] for f in pget(ty, 'fields') if f[0] not in pget(ty, 'anonymous')]
if hasattr(cls, '_anonymous_'):
for anm, aty in [(a, get_aty(a)) for a in cls._anonymous_]:
for fnm in (get_fnms(aty) + flatten([get_fnms(get_aty(aanm, pget(aty, 'fields'))) for aanm in pget(aty, 'anonymous')])):
setattr(cls, fnm, property(lambda self: getattr(getattr(self, anm), fnm), lambda self, v: setattr(getattr(self, anm), fnm, v)))
setattr(cls, '_packed_anonymous_', cls._anonymous_)
setattr(cls, '_anonymous_', [])
type(ctypes.Structure).__setattr__(cls, '_fields_', [('_data', ctypes.c_ubyte * ((offset + 7) // 8))])
type(ctypes.Structure).__setattr__(cls, '_packed_', True)
setattr(cls, '_packed_fields_', fields)
@@ -69,4 +86,3 @@ else:
for f,v in zip(self._packed_fields_, args): setattr(self, f[0], v)
for k,v in kwargs.items(): setattr(self, k, v)
else: super().__init__(*args, **kwargs)
tinygrad/runtime/support/compiler_amd.py
+1 -2
View File
@@ -1,7 +1,6 @@
import ctypes
from tinygrad.helpers import system
import tinygrad.runtime.autogen.comgr as comgr
assert comgr.AMD_COMGR_LANGUAGE_HIP == 4
from tinygrad.runtime.autogen import comgr
try:
comgr.amd_comgr_get_version(ctypes.byref(major:=ctypes.c_uint64()), ctypes.byref(minor:=ctypes.c_uint64()))
if major.value >= 3:
tinygrad/runtime/support/compiler_cpu.py
+5 -5
View File
@@ -1,8 +1,8 @@
import ctypes, platform, sys, subprocess
from tinygrad.device import Compiler
from tinygrad.helpers import OSX, getenv, capstone_flatdump, DEBUG
from tinygrad.helpers import OSX, getenv, capstone_flatdump, DEBUG, unwrap
from tinygrad.runtime.support.elf import jit_loader
try: import tinygrad.runtime.autogen.llvm as llvm
try: from tinygrad.runtime.autogen import llvm
except (ImportError, FileNotFoundError): llvm = None #type:ignore[assignment]
class ClangJITCompiler(Compiler):
@@ -24,7 +24,7 @@ class ClangJITCompiler(Compiler):
def cerr(): return ctypes.pointer(ctypes.pointer(ctypes.c_char()))
def expect(x, err, ret=None):
if x: raise RuntimeError(llvm.string_cast(err.contents) if not isinstance(err, str) else err)
if x: raise RuntimeError(unwrap(ctypes.cast(err.contents, ctypes.c_char_p).value).decode() if not isinstance(err, str) else err)
return ret
class LLVMCompiler(Compiler):
@@ -50,7 +50,7 @@ class LLVMCompiler(Compiler):
self.passes = b'default<O0>'
self.diag_msgs: list[str] = []
@ctypes.CFUNCTYPE(None, llvm.LLVMDiagnosticInfoRef, ctypes.c_void_p)
@llvm.LLVMDiagnosticHandler
def handle_diag(diag_ref, _arg):
severity = llvm.LLVMGetDiagInfoSeverity(diag_ref)
msg = ctypes.string_at(llvm.LLVMGetDiagInfoDescription(diag_ref)).decode()
@@ -70,7 +70,7 @@ class LLVMCompiler(Compiler):
expect(llvm.LLVMRunPasses(mod, self.passes, self.target_machine, self.pbo), 'failed to run passes')
if DEBUG >= 7: print(ctypes.string_at(llvm.LLVMPrintModuleToString(mod)).decode())
obj_buf = expect(llvm.LLVMTargetMachineEmitToMemoryBuffer(self.target_machine, mod, llvm.LLVMObjectFile, err:=cerr(),
ctypes.pointer(buf:=llvm.LLVMMemoryBufferRef())), err, buf)
buf:=llvm.LLVMMemoryBufferRef()), err, buf)
llvm.LLVMDisposeModule(mod)
obj = ctypes.string_at(llvm.LLVMGetBufferStart(obj_buf), llvm.LLVMGetBufferSize(obj_buf))
llvm.LLVMDisposeMemoryBuffer(obj_buf)

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