maybe

* small prg cleanup

* test_timing

.github/workflows/autogen.yml

@@ -19,7 +19,7 @@ on:
 
 jobs:
   autogen:
-    name: Autogen
+    name: In-tree Autogen
     runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
@@ -31,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:
@@ -101,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

@@ -638,13 +638,11 @@ jobs:
     - 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=18 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=11 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

.github/workflows/test.yml

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

autogen_stubs.sh

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

extra/bandwidth_test.py

@@ -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/torch_gemm.py

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

@@ -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/mesa/lvp_nir_options.sh

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

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

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

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

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

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

extra/sqtt/install_sqtt_decoder.py

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

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

@@ -72,7 +72,7 @@ class _ROCParseCtx:
     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]
       inst_execs.append(InstExec(inst_typ, inst_disasm, inst_ev.stall, inst_ev.duration, inst_ev.time))
 
@@ -105,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

extra/sqtt/test_timing.py

@@ -14,7 +14,7 @@ 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[os.environ["DEV"]]
 
@@ -36,7 +36,7 @@ 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
   if os.environ["DEV"] == "AMD":

extra/thunder/tiny/tk/group.py

@@ -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 RT
+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,7 +120,8 @@ 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 = self.ker.alloc((self.group_threads, 2), src.dtype.base, AddrSpace.LOCAL)
@@ -136,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)
@@ -165,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) * RT.TILE_ROW_DIM + (warp_laneid // 4)
-        col = load_i_width * RT.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) * RT.TILE_ROW_DIM + 2 * (warp_laneid % 4)
-        col = load_i_width * RT.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)
 
@@ -211,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) * RT.TILE_ROW_DIM + (warp_laneid // 4)
-      col = store_i_width * RT.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:])
@@ -268,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

@@ -68,10 +68,10 @@ class Kernel(AbstractContextManager):
 
     return uop
 
-  def gl(self, shape, dtype): return GL(shape, dtype, self)._uop
-  def st(self, shape, dtype): return ST(shape, dtype, self)._uop
-  def rt(self, shape, dtype): return RT(shape, dtype, self)._uop
-  def rv(self, length, dtype, layout="naive"): return RV(length, dtype, layout, self)._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))
 
@@ -80,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

@@ -1,36 +1,132 @@
+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
 
+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 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
+
+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__"
+    }
+
+  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__
+
+    for name in dir(source_cls):
+      if name in blacklist or not name.startswith("__"): continue
+
+      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))
+
+    return cls
+  return decorator
+
+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, shape, dtype, ker):
-    self.shape, self.dtype = shape, dtype
-    self._uop = ker.alloc(shape, dtype, AddrSpace.GLOBAL)
+  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, shape, dtype, ker):
-    self.shape, self.dtype = shape, dtype
-    self._uop = ker.alloc(shape, dtype, AddrSpace.LOCAL)
+  def __init__(self, uop, ker):
+    self._uop, self.ker = uop, ker
 
-class RT:
-  TILE_ROW_DIM, TILE_COL_DIM = 16, 16
-  BASE_TILE_NE = TILE_ROW_DIM * TILE_COL_DIM
+  @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, shape, dtype, ker):
+  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.TILE_ROW_DIM == 0
-    assert shape[1] % RT.TILE_COL_DIM == 0
+    assert shape[0] % RT.BASE_TILE_ROWS == 0
+    assert shape[1] % RT.BASE_TILE_COLS == 0
 
-    height = shape[0] // RT.TILE_ROW_DIM
-    width = shape[1] // RT.TILE_COL_DIM
+    height = shape[0] // RT.BASE_TILE_ROWS
+    width = shape[1] // RT.BASE_TILE_COLS
 
-    self.shape, self.dtype = (height, width, self.BASE_TILE_NEPT), dtype
-    self._uop = ker.alloc(self.shape, dtype, AddrSpace.REG)
+    uop = ker.alloc((height, width, RT.BASE_TILE_NEPT), dtype, AddrSpace.REG)
+    return cls(uop, ker)
 
-class RV:
-  def __init__(self, length, dtype, layout, ker):
-    tiles = length // RT.TILE_ROW_DIM
+@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":
@@ -41,5 +137,7 @@ class RV:
         outer_dim = tiles
       case _: raise NotImplementedError(f"rv layout {layout} not implemented")
 
-    self.shape, self.dtype = (outer_dim, inner_dim, 2), dtype
-    self._uop = ker.alloc(self.shape, dtype, AddrSpace.REG)
+    uop = ker.alloc((outer_dim, inner_dim, 2), dtype, AddrSpace.REG)
+    return RV(uop, ker)
+
+ALL_TILES = UOp | GL | ST | RT | RV

pyproject.toml

@@ -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/mockgpu/amd/amdgpu.py

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

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

@@ -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,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_tiny.py

@@ -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/testextra/test_tk.py

@@ -152,6 +152,44 @@ 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
@@ -365,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()
 
@@ -379,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)

test/unit/test_autogen.py

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

tinygrad/helpers.py

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

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

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

@@ -4,19 +4,26 @@ from typing import TypeAlias, TYPE_CHECKING, Self
 from tinygrad.uop import Ops
 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
+
+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
@@ -42,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)
@@ -85,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`.
@@ -110,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:
@@ -129,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:
@@ -150,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
 
@@ -163,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.
@@ -179,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`.
 
@@ -234,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.
 
@@ -251,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
 
@@ -264,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.
 
@@ -306,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:
     """
@@ -323,28 +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:
+  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_)]
+    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_)]
+    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_)])
+    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_)]).reshape(noop + flatten((k,(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)))
+    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_)))
+    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_))])
+    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/state.py

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

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

@@ -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/am/__init__.py

@@ -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/nvjitlink.py

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

@@ -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/rocprof.py

@@ -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/ops_amd.py

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

tinygrad/runtime/ops_cuda.py

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

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

tinygrad/runtime/ops_qcom.py

@@ -321,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)
@@ -357,8 +359,8 @@ class QCOMDevice(HCQCompiled):
                      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)

tinygrad/runtime/ops_webgpu.py

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

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

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

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

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

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

tinygrad/runtime/support/compiler_cuda.py

@@ -1,7 +1,7 @@
 import subprocess, hashlib, tempfile, ctypes, re, pathlib
 from typing import Callable
 from tinygrad.helpers import to_char_p_p, colored, init_c_var, getenv, system
-import tinygrad.runtime.autogen.nvrtc as nvrtc
+from tinygrad.runtime.autogen import nvrtc, nvjitlink as jitlink
 from tinygrad.device import Compiler, CompileError
 
 CUDA_PATH = getenv("CUDA_PATH", "")
@@ -17,8 +17,8 @@ def nvrtc_check(status, ctx=None):
 
 def jitlink_check(status, ctx=None):
   if status != 0:
-    err_log = _get_bytes(ctx, nvrtc.nvJitLinkGetErrorLog, nvrtc.nvJitLinkGetErrorLogSize, lambda _: None).decode() if ctx else ""
-    raise CompileError(f"NvJitLink Error {status}, {nvrtc.nvJitLinkResult__enumvalues.get(status, 'Unknown')}\n{err_log}")
+    err_log = _get_bytes(ctx, jitlink.nvJitLinkGetErrorLog, jitlink.nvJitLinkGetErrorLogSize, lambda _: None).decode() if ctx else ""
+    raise CompileError(f"jitlink Error {status}, {jitlink.nvJitLinkResult.get(status)}\n{err_log}")
 
 def pretty_ptx(s):
   # all expressions match `<valid_before><expr><valid_after>` and replace it with `<valid_before>color(<expr>)<valid_after>`
@@ -83,12 +83,12 @@ class PTXCompiler(Compiler):
 
 class NVPTXCompiler(PTXCompiler):
   def __init__(self, arch:str):
-    nvrtc_check(nvrtc.nvJitLinkVersion(ctypes.byref(ctypes.c_uint()), ctypes.byref(ctypes.c_uint())))
+    nvrtc_check(jitlink.nvJitLinkVersion(ctypes.byref(ctypes.c_uint()), ctypes.byref(ctypes.c_uint())))
     super().__init__(arch, cache_key="nv_ptx")
   def compile(self, src:str) -> bytes:
-    jitlink_check(nvrtc.nvJitLinkCreate(handle := nvrtc.nvJitLinkHandle(), 1, to_char_p_p([f'-arch={self.arch}'.encode()])), handle)
-    jitlink_check(nvrtc.nvJitLinkAddData(handle, nvrtc.NVJITLINK_INPUT_PTX, ptxsrc:=super().compile(src), len(ptxsrc), "<null>".encode()), handle)
-    jitlink_check(nvrtc.nvJitLinkComplete(handle), handle)
-    data = _get_bytes(handle, nvrtc.nvJitLinkGetLinkedCubin, nvrtc.nvJitLinkGetLinkedCubinSize, jitlink_check)
-    jitlink_check(nvrtc.nvJitLinkDestroy(handle))
+    jitlink_check(jitlink.nvJitLinkCreate(handle := jitlink.nvJitLinkHandle(), 1, to_char_p_p([f'-arch={self.arch}'.encode()])), handle)
+    jitlink_check(jitlink.nvJitLinkAddData(handle, jitlink.NVJITLINK_INPUT_PTX, ptxsrc:=super().compile(src), len(ptxsrc), "<null>".encode()), handle)
+    jitlink_check(jitlink.nvJitLinkComplete(handle), handle)
+    data = _get_bytes(handle, jitlink.nvJitLinkGetLinkedCubin, jitlink.nvJitLinkGetLinkedCubinSize, jitlink_check)
+    jitlink_check(jitlink.nvJitLinkDestroy(handle))
     return data

tinygrad/runtime/support/compiler_mesa.py

@@ -1,9 +1,9 @@
 import base64, ctypes, pathlib, tempfile, hashlib
 from tinygrad.device import Compiler
 from tinygrad.helpers import cpu_objdump, system
-import tinygrad.runtime.autogen.mesa as mesa
+from tinygrad.runtime.autogen import mesa
 from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, expect, cerr
-try: import tinygrad.runtime.autogen.llvm as llvm
+try: from tinygrad.runtime.autogen import llvm
 except (ImportError, FileNotFoundError): llvm = None #type:ignore[assignment]
 
 def deserialize(enc_src, opts):

tinygrad/runtime/support/ib.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
-import resource, ctypes, weakref, functools, itertools, tinygrad.runtime.autogen.ib as ib
+import resource, ctypes, weakref, functools, itertools
+from tinygrad.runtime.autogen import ib
 from typing import Iterator
 from dataclasses import dataclass
 from weakref import WeakKeyDictionary
@@ -141,7 +142,7 @@ class IBConn:
     while (wr_id in self.pending_wrids) if wr_id is not None else self.pending_wrids:
       if self.ctx.ctx.contents.ops.poll_cq(self.cq, _num_entries:=1, ctypes.byref(wc:=ib.struct_ibv_wc())):
         if wc.status != ib.IBV_WC_SUCCESS:
-          raise RuntimeError(f'Work Request completed with error: wr_id={wc.wr_id} status={ib.ibv_wc_status__enumvalues.get(wc.status, wc.status)}')
+          raise RuntimeError(f'Work Request completed with error: wr_id={wc.wr_id} status={ib.enum_ibv_wc_status.get(wc.status, wc.status)}')
         self.pending_wrids.remove(wc.wr_id)
 
   def rdma_write(self, sgl:list[SGE]):
@@ -162,7 +163,7 @@ class IBConn:
         # Scatter-Gather Entry for local memory
         sge = ctypes.pointer(ib.struct_ibv_sge(addr=sg.src_iova+off, length=min(sg.size-off, self.ctx.port_attr.max_msg_sz), lkey=sg.src_key))
         # RDMA struct for remote memory
-        wr = ib.union_ibv_send_wr_wr(rdma=ib.struct_ibv_send_wr_1_rdma(remote_addr=sg.dst_iova+off, rkey=sg.dst_key))
+        wr = ib.struct_ibv_send_wr_wr(rdma=ib.struct_ibv_send_wr_wr_rdma(remote_addr=sg.dst_iova+off, rkey=sg.dst_key))
         # Signal (with chosen work request id) if it's the last wr (first in the loop since it's reversed)
         wid, flags = (wr_id, ib.IBV_SEND_SIGNALED) if swr is None else (0, 0)
         # Create Send Request

tinygrad/runtime/support/nv/ip.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 import ctypes, time, array, struct, itertools, dataclasses
 from typing import cast, Any
-from tinygrad.runtime.autogen.nv import nv
+from tinygrad.runtime.autogen import nv, nv_570 as nv_gpu, pci
 from tinygrad.helpers import to_mv, lo32, hi32, DEBUG, round_up, round_down, mv_address, fetch, wait_cond
 from tinygrad.runtime.support.system import System
 from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.autogen import nv_gpu, pci
 
 @dataclasses.dataclass(frozen=True)
 class GRBufDesc: size:int; virt:bool; phys:bool; local:bool=False # noqa: E702
@@ -69,8 +68,8 @@ class NVRpcQueue:
       System.memory_barrier()
 
       if DEBUG >= 3:
-        rpc_names = {**nv.c__Ea_NV_VGPU_MSG_FUNCTION_NOP__enumvalues, **nv.c__Ea_NV_VGPU_MSG_EVENT_FIRST_EVENT__enumvalues}
-        print(f"nv {self.gsp.nvdev.devfmt}: in RPC: {rpc_names.get(hdr.function, f'ev:{hdr.function:x}')}, res:{hdr.rpc_result:#x}")
+        nm = nv.rpc_fns.get(hdr.function, nv.rpc_events.get(hdr.function, f'ev:{hdr.function:x}'))
+        print(f"nv {self.gsp.nvdev.devfmt}: in RPC: {nm}, res:{hdr.rpc_result:#x}")
 
       if hdr.rpc_result != 0: raise RuntimeError(f"RPC call {hdr.function} failed with result {hdr.rpc_result}")
       if hdr.function == cmd: return msg
@@ -443,7 +442,7 @@ class NV_GSP(NV_IP):
     bufs_p = nv_gpu.struct_NV90F1_CTRL_VASPACE_COPY_SERVER_RESERVED_PDES_PARAMS(pageSize=res_sz, numLevelsToCopy=3,
       virtAddrLo=res_va, virtAddrHi=res_va + res_sz - 1)
     for i,pt in enumerate(self.nvdev.mm.page_tables(res_va, size=res_sz)):
-      bufs_p.levels[i] = nv_gpu.struct_NV90F1_CTRL_VASPACE_COPY_SERVER_RESERVED_PDES_PARAMS_0(physAddress=pt.paddr,
+      bufs_p.levels[i] = nv_gpu.struct_NV90F1_CTRL_VASPACE_COPY_SERVER_RESERVED_PDES_PARAMS_level(physAddress=pt.paddr,
         size=self.nvdev.mm.pte_cnt[0] * 8 if i == 0 else 0x1000, pageShift=self.nvdev.mm.pte_covers[i].bit_length() - 1, aperture=1)
     self.rpc_rm_control(hObject=vaspace, cmd=nv_gpu.NV90F1_CTRL_CMD_VASPACE_COPY_SERVER_RESERVED_PDES, params=bufs_p)
 

tinygrad/schedule/indexing.py

@@ -6,7 +6,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, g
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
 
-ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
                      Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL}
 
@@ -16,10 +16,9 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
   for s in rb.src:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
-def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
+def realize_store(ctx:dict[UOp, None], a:UOp) -> None:
   if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
-  # if it's a kernel, we don't realize it
-  if a.src[1].op is not Ops.KERNEL: ctx[a] = None
+  ctx[a] = None
 
 pm_generate_realize_map = PatternMatcher([
   # always realize SINK src
@@ -29,7 +28,7 @@ pm_generate_realize_map = PatternMatcher([
   # realize srcs of COPY, MSELECT, MSTACK
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
   # realize ASSIGN and input to assign (might be optimized out)
-  (UPat(Ops.ASSIGN, name="a"), realize_assign),
+  (UPat(Ops.STORE, name="a"), realize_store),
 ])
 
 @dataclass(frozen=True)
@@ -56,16 +55,21 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
   new_srcs = []
   for s in x.src:
     new_src = s
-    if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.AFTER and s.src[1].op is Ops.KERNEL):
+    if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
       if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
     elif s in ctx.realize_map:
       realized_ranges = ctx.realize_map[s]
       assert isinstance(realized_ranges, list), "realize map must contain range list"
       closed_ranges = tuple([r for i,r in enumerate(ctx.range_map[s][1]) if i in realized_ranges])
-      # None in the device assigns it a number later
-      opts = BufferizeOpts(device=s.device) if len(ctx.range_map[s][1]) == len(realized_ranges) else BufferizeOpts(None, AddrSpace.LOCAL)
-      new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+closed_ranges, arg=opts, tag=s.tag if opts.addrspace == AddrSpace.GLOBAL else None)
-      if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][0]) if i in realized_ranges])
+      if s.op is Ops.STORE:
+        # add the ends if this is a store
+        new_src = s.end(*[r for r in closed_ranges if r.op is Ops.RANGE])
+        del ctx.realize_map[s]
+      else:
+        # None in the device assigns it a number later
+        opts = BufferizeOpts(device=s.device) if len(ctx.range_map[s][1]) == len(realized_ranges) else BufferizeOpts(None, AddrSpace.LOCAL)
+        new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+closed_ranges, arg=opts, tag=s.tag if opts.addrspace == AddrSpace.GLOBAL else None)
+        if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][0]) if i in realized_ranges])
     new_srcs.append(new_src)
   # NOTE: do we need this?
   return x.replace(src=tns) if x.src != (tns:=tuple(new_srcs)) else None
@@ -85,7 +89,7 @@ def convert_reduce_axis_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
   return ret
 
 def remove_movement_op_after_rangeify(ctx:IndexingContext, x:UOp):
-  if x in ctx.range_map or x.src[0].op is Ops.INDEX: return x.src[0]
+  if (x in ctx.range_map or x.src[0].op is Ops.INDEX): return x.src[0]
 
 def add_third_op_to_assign_to_track_shape(ctx:IndexingContext, assign:UOp):
   if assign.src[1].op is Ops.KERNEL: return None
@@ -176,7 +180,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       # mark all ranges as ended
       assert rctx.realize_map[x] is None
       rctx.realize_map[x] = list(range(len(x.shape)))
-    elif x.op in {Ops.MSTACK, Ops.MSELECT}:
+    elif x.op in {Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
       # treat MSTACK/MSELECT like SINK
       continue
     elif len(consumer_rngs) == 0:
@@ -209,6 +213,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       out_rngs = tuple(_out_rngs)
 
       # we have to (partially) realize here if there's new ranges
+      print(_realize_axis)
       if len(_realize_axis): rctx.realize_map[x] = _realize_axis
 
     # if this element is a reduce and there's ended ranges, we might have to end some other ranges

tinygrad/schedule/rangeify.py

@@ -534,6 +534,12 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
 
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
 
+  # link any movementops with tags to sink, and remove the tags from other parts of the graph
+  # we add "None" to the tag so it's deduped from the movementop that exists in the graph
+  #tagged_mops = [x.replace(tag=x.tag+(None,)) for x in tsink.toposort() if x.op in GroupOp.Movement and x.tag is not None]
+  #tsink = tsink.replace(src=tsink.src+tuple(tagged_mops))
+  #tsink = tsink.substitute({})
+
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, DEBUG_RANGEIFY)
 

tinygrad/tensor.py

@@ -299,7 +299,7 @@ class Tensor(OpMixin):
     assert self.shape == x.shape, f"assign shape mismatch {self.shape} != {x.shape}"
     assert self.device == x.device, f"assign device mismatch {self.device} != {x.device}"
     assert self.dtype == x.dtype, f"assign dtype mismatch {self.dtype} != {x.dtype}"
-    return self.replace(self._apply_uop(UOp.assign, x))
+    return self.replace(self._apply_uop(lambda x,y: x.after(x.store(y)), x))
 
   def detach(self) -> Tensor:
     """

tinygrad/viz/serve.py

@@ -215,9 +215,9 @@ def load_sqtt(profile:list[ProfileEvent]) -> None:
   if not rctx.inst_execs: return err("EMPTY SQTT OUTPUT", f"{len(sqtt_events)} SQTT events recorded, none got decoded")
   steps:list[dict] = []
   for name,waves in rctx.inst_execs.items():
-    if (r:=ref_map.get(name)): name = ctxs[r]["name"]
-    steps.append({"name":name, "depth":0, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
-                  "data":{"src":trace.keys[r].ret.src if r else name, "lang":"cpp"}})
+    prg = trace.keys[r].ret if (r:=ref_map.get(name)) else None
+    steps.append({"name":prg.name if prg is not None else name, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
+                  "depth":0, "data":{"src":prg.src if prg is not None else name, "lang":"cpp"}})
 
     # Idle:     The total time gap between the completion of previous instruction and the beginning of the current instruction.
     #           The idle time can be caused by: