grad outerworld test

with rangify we can write these with arange

.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,67 +31,113 @@ jobs:
         opencl: 'true'
         amd: 'true'
         cuda: 'true'
-        webgpu: 'true'
         llvm: 'true'
+        webgpu: 'true'
+        mesa: 'true'
         pydeps: '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
-    runs-on: ubuntu-24.04
+    - name: Verify libclang autogen
+      run: |
+        cp tinygrad/runtime/autogen/libclang.py /tmp/libclang.py.bak
+        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
+        diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
+  autogen-mac:
+    name: In-tree Autogen (macos)
+    runs-on: macos-14
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -99,11 +145,32 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       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
+        llvm: 'true'
+    - name: Verify macos 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
+        mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
+        LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
+        diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
+  autogen-comgr-3:
+    name: In-tree Autogen (comgr 3)
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+    - name: Install autogen support packages
+      run: |
+        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
+        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
+        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.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
@@ -287,8 +287,8 @@ jobs:
         python extra/optimization/extract_dataset.py
         gzip -c /tmp/sops > extra/datasets/sops.gz
         #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 18500 lines
-      run: MAX_LINE_COUNT=18500 python sz.py
+    - name: Repo line count < 19000 lines
+      run: MAX_LINE_COUNT=19000 python sz.py
 
   spec:
     strategy:

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,40 @@ 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):
+            base_row = (local_warpid * dst.shape[-3] + height) * RT.BASE_TILE_ROWS
+            base_col = width * RT.BASE_TILE_COLS
 
-        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)
+            if not transpose:
+              row = base_row + (warp_laneid // 4)
+              col = base_col + 2 * (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 % 4) // 2) * 8
+              col_offset = (inner % 2) + (inner // 4) * 8
+            else:
+              row = base_row + 2 * (warp_laneid % 4)
+              col = base_col + (warp_laneid // 4)
 
-      src_i_last = (row + row_offset) * src.shape[-1] + col + col_offset
+              row_offset = (inner % 2) + (inner // 4) * 8
+              col_offset = ((inner % 4) // 2) * 8
 
-      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)
+            src_i_last = (row + row_offset) * src.shape[-1] + col + col_offset
+
+            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 +210,56 @@ 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: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 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):
+            base_row = (local_warpid * src.shape[-3] + height) * RT.BASE_TILE_ROWS
+            base_col = width * RT.BASE_TILE_COLS
 
-      row_offset = ((store_i_inner % 4) // 2) * 8
-      col_offset = (store_i_inner % 2) + (store_i_inner // 4) * 8
+            if not transpose:
+              row = base_row + (warp_laneid // 4)
+              col = base_col + 2 * (warp_laneid % 4)
 
-      dst_i_last = (row + row_offset) * dst.shape[-1] + col + col_offset
+              row_offset = ((inner % 4) // 2) * 8
+              col_offset = (inner % 2) + (inner // 4) * 8
+            else:
+              row = base_row + 2 * (warp_laneid % 4)
+              col = base_col + (warp_laneid // 4)
 
-      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)
+              row_offset = (inner % 2) + (inner // 4) * 8
+              col_offset = ((inner % 4) // 2) * 8
+
+            dst_i_last = (row + row_offset) * dst.shape[-1] + col + col_offset
+
+            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,20 +273,18 @@ 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")
 
     self.ker.push_store(dst_store, dst)
-    return dst.after(dst_store.barrier()).reshape(dst.shape) if after else dst_store
+    return dst.after(dst_store.barrier()).reshape(dst.shape)

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/external/external_openpilot_image_warp.py

@@ -0,0 +1,55 @@
+import time
+from tinygrad.tensor import Tensor, Device
+
+MODEL_WIDTH = 512
+MODEL_HEIGHT = 256
+MODEL_FRAME_SIZE = MODEL_WIDTH * MODEL_HEIGHT * 3 // 2
+IMG_INPUT_SHAPE = (1, 12, 128, 256)
+
+def tensor_arange(end): return Tensor([float(i) for i in range(end)])
+def tensor_round(tensor:Tensor): return (tensor + 0.5).floor()
+
+h_src, w_src = 1208, 1928
+h_dst, w_dst = MODEL_HEIGHT, MODEL_WIDTH
+x = tensor_arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst)
+y = tensor_arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst)
+ones = Tensor.ones_like(x)
+dst_coords = x.reshape((1,-1)).cat(y.reshape((1,-1))).cat(ones.reshape((1,-1)))
+
+def warp_perspective_tinygrad(src:Tensor, M_inv:Tensor) -> Tensor:
+  src_coords = M_inv @ dst_coords
+  src_coords = src_coords / src_coords[2:3, :]
+
+  x_src = src_coords[0].reshape(h_dst, w_dst)
+  y_src = src_coords[1].reshape(h_dst, w_dst)
+
+  x_nearest = tensor_round(x_src).clip(0, w_src - 1).cast('int')
+  y_nearest = tensor_round(y_src).clip(0, h_src - 1).cast('int')
+
+  # TODO: make 2d indexing fast
+  idx = y_nearest*src.shape[1] + x_nearest
+  dst = src.flatten()[idx]
+  return dst.reshape(h_dst, w_dst)
+
+if __name__ == "__main__":
+  from tinygrad.engine.jit import TinyJit
+  update_img_jit = TinyJit(warp_perspective_tinygrad, prune=True)
+
+  step_times = []
+  for _ in range(10):
+    # regenerate inputs
+    inputs = [Tensor.randn(1928,1208), Tensor.randn(3,3)]
+    Tensor.realize(*inputs)
+    Device.default.synchronize()
+
+    # do the warp
+    st = time.perf_counter()
+    out = update_img_jit(*inputs)
+    mt = time.perf_counter()
+    val = out.contiguous().realize()
+    Device.default.synchronize()
+    et = time.perf_counter()
+
+    # measure the time
+    step_times.append((et-st)*1e3)
+    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")

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_arange.py

@@ -7,24 +7,28 @@ from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.uop.ops import Ops
 
 class TestArange(unittest.TestCase):
-  def _get_flops(self, N):
+  def _get_flops(self, tensor, desired):
     GlobalCounters.reset()
-    tt = Tensor.arange(N)
-    sched = tt.schedule()
+    sched = tensor.schedule()
     self.assertEqual(len(sched), 1)
     p = get_program(sched[-1].ast)
-    ExecItem(CompiledRunner(p), [tt.uop.buffer]).run()
-    np.testing.assert_equal(tt.numpy(), np.arange(N))
+    ExecItem(CompiledRunner(p), [tensor.uop.buffer]).run()
+    np.testing.assert_equal(tensor.numpy(), desired)
     return p.estimates.ops
 
-  def test_complexity(self):
-    self.assertEqual(self._get_flops(256), 0)
-    self.assertEqual(self._get_flops(2560), 0)
+  def test_arange_complexity(self):
+    self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0)
+    self.assertEqual(self._get_flops(Tensor.arange(2560), np.arange(2560)), 0)
 
   def test_arange_cat(self):
     t = Tensor.arange(2, dtype=dtypes.int)+Tensor([3])
     self.assertEqual(t.cat(t).tolist(), [3, 4, 3, 4])
 
+  def test_eye_complexity(self):
+    with Context(NOOPT=1):
+      # NOTE: not every backend supports CMPEQ
+      self.assertLessEqual(self._get_flops(Tensor.eye(2560).contiguous(), np.eye(2560)), 2*2560*2560)
+
 DSET, DDIM = 2048, 32
 
 class TestIndexing(unittest.TestCase):

test/test_linearizer.py

@@ -14,6 +14,8 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
 MOCKGPU = getenv("MOCKGPU")
 
+from tinygrad.uop.ops import print_uops # noqa: F401 # pylint: disable=unused-import
+
 class TestLinearizer(unittest.TestCase):
   def test_arg_dedup(self):
     # NOTE: this realize exists because Tensor.numpy calls .contiguous() internally
@@ -38,6 +40,22 @@ class TestLinearizer(unittest.TestCase):
     np.testing.assert_equal(a.numpy(), ta)
     np.testing.assert_equal(b.numpy(), tb)
 
+  @unittest.skip("TODO: some backends insert more casts")
+  def test_cast_there_and_back(self):
+    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
+    out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
+    ast = helper_linearizer_opt(out)
+    uops = get_program(ast, opts=[]).uops
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
+
+  @unittest.expectedFailure
+  def test_cast_back_and_there(self):
+    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
+    out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
+    ast = helper_linearizer_opt(out)
+    uops = get_program(ast, opts=[]).uops
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
+
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
   def test_late_bias_load(self):
     img = Tensor.empty(1, 3, 16, 16)
@@ -491,7 +509,7 @@ def copyout_outputs(outbufs:list[Buffer]) -> list[np.ndarray]:
   return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
 
 def reset_bufs(bufs:list[Buffer]):
-  for buf in bufs: buf.copyin(np.zeros((buf.size, ), dtype=_to_np_dtype(buf.dtype)).data) # Zero to check that all values are filled
+  for buf in bufs: buf.copyin(np.zeros((buf.size*buf.dtype.itemsize,), dtype=np.uint8).data)
 
 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
                                apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):

test/test_outerworld.py

@@ -11,6 +11,87 @@ class TestOuterworldReduce(unittest.TestCase):
     t = Tensor(UOp(Ops.REDUCE, dtype=out.uop.dtype, src=(out.uop, a), arg=Ops.ADD))
     self.assertListEqual(t.tolist(), [5.,5.,5.,5.,5.])
 
+# TODO: delete test_outerworld_range?
+class TestOuterRange(unittest.TestCase):
+  def test_simple_range(self):
+    a = Tensor.ones(10).contiguous()
+    acc = Tensor.zeros().contiguous()
+    Tensor.realize(a, acc)
+
+    # this is fold
+    i = UOp.range(10, -100, AxisType.OUTER)
+    acc_i = acc.uop.after(i)
+    vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
+    out = Tensor(acc.uop.after(acc_i.store(acc_i + a[vi].uop).end(i)))
+    out.realize()
+    assert out.item() == 10.0
+
+  def test_inner_range(self):
+    a = Tensor.ones(10, 10).contiguous()
+    acc = Tensor.zeros(10).contiguous()
+    Tensor.realize(a, acc)
+
+    # this is fold
+    i = UOp.range(10, -100, AxisType.OUTER)
+    acc_i = acc.uop.after(i)
+    vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
+    out = Tensor(acc.uop.after(acc_i.store(acc_i + a[:, vi].uop).end(i)))
+    out.realize()
+    assert all(x == 10.0 for x in out.tolist())
+
+  def test_range_matmul(self):
+    vec = Tensor.randn(1, 10).realize()
+    mats = Tensor.randn(3, 10, 10).realize()
+
+    # 3 matmuls in "scan"
+    ref = ((vec @ mats[0]) @ mats[1]) @ mats[2]
+    ref.realize()
+
+    # 3 matmuls with outer world range
+    i = UOp.range(3, -100, AxisType.OUTER)
+    vec_i = Tensor(vec.uop.after(i))
+    vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
+    out = Tensor(vec.uop.after(vec_i.uop.store((vec_i.contiguous() @ mats[vi]).uop).end(i)))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+
+  def test_range_grad(self):
+    def range_matmul(vec, mats):
+      # vec: (1, 10), mats: (3, 10, 10)
+      # assume vec, mats already have requires_grad set however you like
+
+      i = UOp.range(3, -100, AxisType.OUTER)      # loop axis
+      vec_i = Tensor(vec.uop.after(i))            # "loop-carried" vector
+      vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
+
+      body = (vec_i.contiguous() @ mats[vi])      # matmul using loop index
+      out = Tensor(vec.uop.after(vec_i.uop.store(body.uop).end(i)))
+      return out
+
+    vec = Tensor.randn(1, 3, requires_grad=True)
+    mats = Tensor.randn(3, 3, 3, requires_grad=True)
+    Tensor.realize(vec, mats)
+
+    ref = ((vec @ mats[0]) @ mats[1]) @ mats[2]
+    loss = (1.0 - ref).square().mean()
+    loss.backward()
+    Tensor.realize(vec.grad, mats.grad)
+    print(vec.grad.numpy())
+    print(mats.grad.numpy())
+    vec.grad = None
+    mats.grad = None
+
+    out = range_matmul(vec, mats)
+    loss = (1.0 - out).square().mean()
+    loss.backward()
+    Tensor.realize(vec.grad, mats.grad)
+
+    print(vec.grad, mats.grad)   # should be non-None and finite
+    print(vec.grad.numpy())
+    print(mats.grad.numpy())
+
 class TestOuterworld(unittest.TestCase):
   def test_range_plus_1(self):
     t = Tensor.arange(100).reshape(10,10).realize()

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/codegen/opt/postrange.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 import math, itertools
 from collections import defaultdict
 from typing import cast, Final
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp, axis_letters, axis_colors
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp
+from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
@@ -12,10 +13,6 @@ from tinygrad.renderer import Renderer
 
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
-# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
-               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
-
 class Scheduler:
   def __init__(self, ast:UOp, ren:Renderer):
     self.ast, self.ren = ast, ren

tinygrad/engine/schedule.py

@@ -1,5 +1,5 @@
 from typing import cast
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from collections import deque, defaultdict
 from tinygrad.uop.ops import UOp, Ops, buffers
 from tinygrad.device import Device, Buffer, MultiBuffer
@@ -13,6 +13,7 @@ class ScheduleItem:
   bufs: tuple[Buffer, ...]
   metadata: tuple[Metadata, ...] = ()
   fixedvars: dict[str, int] = field(default_factory=dict)
+  bound_ranges: tuple[UOp, ...] = ()
 
 # **** schedule linearizer
 
@@ -25,7 +26,8 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
     if u.op is not Ops.AFTER: continue  # anything that's not an ASSIGN doesn't write a kernel, so we can skip
     k = u.src[1]
     in_degree.setdefault(k, 0)
-    for s in k.src:
+    if k.op is Ops.RANGE: continue
+    for s in k.src[0].src if k.op is Ops.END else k.src:
       if s.op is Ops.AFTER:
         children[s.src[1]].append(k)
         in_degree[k] += 1
@@ -39,16 +41,19 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
       elif s.op is Ops.BUFFER:
         pass  # a BUFFER is already realized, nothing to do here
       elif s.op is Ops.BIND:
-        var, val = s.unbind()
-        assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
-        var_vals[var.expr] = val
+        # for RANGE this is in fixedvars
+        if s.src[1].op is not Ops.RANGE:
+          var, val = s.unbind()
+          assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
+          var_vals[var.expr] = val
       else:
         raise RuntimeError(f"input to kernel must be AFTER or BUFFER, not {s.op}")
 
   # linearize KERNEL UOps into ScheduleItems in BFS order
 
   def _heuristic(k: UOp):
-    if k.arg.ast.op is Ops.COPY and not all_same([Device[cast(Buffer, s.buf_uop.buffer).device].group_id for s in k.src]): return 1000
+    if k.op is Ops.KERNEL and k.arg.ast.op is Ops.COPY and not all_same([Device[cast(Buffer, s.buf_uop.buffer).device].group_id for s in k.src]):
+      return 1000
     return 0
 
   last_heuristic: int = 0
@@ -57,27 +62,53 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
   for k,v in in_degree.items():
     if v == 0: queues[_heuristic(k)].append(k)
 
-  schedule: list[ScheduleItem] = []
+  schedule: list[ScheduleItem|UOp] = []
   while last_queue or any(queues.values()):
     if not last_queue: last_heuristic, last_queue = min((it for it in queues.items() if it[1]), key=lambda x: abs(x[0]-last_heuristic))
-    k = last_queue.popleft()
-    ast = k.arg.ast
-    # create subbuffers if needed
-    if ast.op is Ops.BUFFER_VIEW:
-      base = k.src[1].buf_uop.buffer
-      assert isinstance(base, Buffer), "base can't be MultiBuffer"
-      buffers[k.src[0]] = base.view(k.size, ast.dtype, ast.arg[1]*base.dtype.itemsize)
-    ubufs = tuple(s.buf_uop.buffer for s in k.src if s.op is not Ops.BIND)
-    if any(isinstance(x, MultiBuffer) for x in ubufs):
-      assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
-      dnums = [x for x in ast.variables() if x.arg[0] == '_device_num']
-      for i,bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
-        schedule.append(ScheduleItem(ast, bufs, k.arg.metadata, {dnums[0].expr:i} if len(dnums) else {}))
+    k = rk = last_queue.popleft()
+    if k.op is Ops.END: k = k.src[0]
+    if k.op is Ops.RANGE: schedule.append(k)
+    elif k.op is Ops.KERNEL:
+      ast = k.arg.ast
+      # create subbuffers if needed
+      if ast.op is Ops.BUFFER_VIEW:
+        base = k.src[1].buf_uop.buffer
+        assert isinstance(base, Buffer), "base can't be MultiBuffer"
+        buffers[k.src[0]] = base.view(k.size, ast.dtype, ast.arg[1]*base.dtype.itemsize)
+      ubufs = tuple(s.buf_uop.buffer for s in k.src if s.op is not Ops.BIND)
+      bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and s.src[1].op is Ops.RANGE)
+      if any(isinstance(x, MultiBuffer) for x in ubufs):
+        assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
+        dnums = [x for x in ast.variables() if x.arg[0] == '_device_num']
+        for i,bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
+          schedule.append(ScheduleItem(ast, bufs, k.arg.metadata, {dnums[0].expr:i} if len(dnums) else {}, bound_ranges=bound_ranges))
+      else:
+        # ONE -> ONE
+        schedule.append(ScheduleItem(ast, cast(tuple[Buffer, ...], ubufs), k.arg.metadata, bound_ranges=bound_ranges))
+      if rk.op is Ops.END: schedule.append(rk)
     else:
-      # ONE -> ONE
-      schedule.append(ScheduleItem(ast, cast(tuple[Buffer, ...], ubufs), k.arg.metadata))
+      raise RuntimeError(f"can't schedule {k.op}")
     for x in children[k]:
       in_degree[x] -= 1
       if in_degree[x] == 0: queues[_heuristic(x)].append(x)
 
-  return schedule, var_vals
+  # expand the ranges in the schedule
+  real_schedule: list[ScheduleItem] = []
+  sched_ptr = 0
+  in_ranges = {}
+  range_ptrs = {}
+  while sched_ptr < len(schedule):
+    si = schedule[sched_ptr]
+    if isinstance(si, UOp):
+      if si.op is Ops.RANGE:
+        in_ranges[si] = 0
+        range_ptrs[si] = sched_ptr + 1
+      elif si.op is Ops.END:
+        if in_ranges[si.src[1]] < si.src[1].vmax:
+          in_ranges[si.src[1]] += 1
+          sched_ptr = range_ptrs[si.src[1]]
+          continue
+    else:
+      real_schedule.append(replace(si, fixedvars=si.fixedvars | {s.src[0].arg[0]:in_ranges[s.src[1]] for s in si.bound_ranges}, bound_ranges=()))
+    sched_ptr += 1
+  return real_schedule, var_vals

tinygrad/gradient.py

@@ -43,6 +43,17 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.KERNEL, name="k"), lambda ctx, k: k.arg.grad_fxn(ctx, k)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),
+
+  # RANGE: loop index / axis, not a differentiable quantity
+  (UPat(Ops.RANGE), lambda: (None,)),
+
+  # STORE: buffer write. Gradient flows only into the value being stored.
+  # src layout is roughly (buffer, value, *axes_or_indices)
+  (UPat(Ops.STORE), lambda ctx: (None, ctx)),
+
+  # END: loop terminator / "end of range" node.
+  # Just pass the gradient into the body (first src), ignore the ranges.
+  (UPat(Ops.END, name="ret"), lambda ctx, ret: (ctx, *[None]*(len(ret.src) - 1))),
 ])
 
 def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]:

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,24 @@
-import importlib, pathlib
-from tinygrad.helpers import system
+import glob, importlib, pathlib, subprocess, tarfile
+from tinygrad.helpers import fetch, flatten, system, getenv
 
 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"}
+macossdk = "/var/db/xcode_select_link/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk"
 
 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
+  if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists() or getenv('REGEN'):
+    files, kwargs['args'] = files() if callable(files) else files, args() if callable(args:=kwargs.get('args', [])) else args
+    if (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 +27,106 @@ 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 "libclang":
+      return load("libclang", ["os.getenv('LIBCLANG_PATH', find_library('clang-20'))"],
+                  lambda: [f"{system('llvm-config-20 --includedir')}/clang-c/{s}.h" for s in ["Index", "CXString", "CXSourceLocation", "CXFile"]],
+                  args=lambda: system("llvm-config-20 --cflags").split())
+    case "metal":
+      return load("metal", ["find_library('Metal')"],[f"{macossdk}/System/Library/Frameworks/Metal.framework/Headers/MTL{s}.h" for s in
+                  ["ComputeCommandEncoder", "ComputePipeline", "CommandQueue", "Device", "IndirectCommandBuffer", "Resource", "CommandEncoder"]],
+                  args=["-xobjective-c","-isysroot",macossdk], types={"dispatch_data_t":"objc.id_"})
     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/libc.py

@@ -13,115 +13,87 @@ off_t = ctypes.c_int64
 mode_t = ctypes.c_uint32
 size_t = ctypes.c_uint64
 __off_t = ctypes.c_int64
-# extern void *mmap(void *__addr, size_t __len, int __prot, int __flags, int __fd, __off_t __offset) __attribute__((nothrow))
 try: (mmap:=dll.mmap).restype, mmap.argtypes = ctypes.c_void_p, [ctypes.c_void_p, size_t, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_int64]
 except AttributeError: pass
 
-# extern int munmap(void *__addr, size_t __len) __attribute__((nothrow))
 try: (munmap:=dll.munmap).restype, munmap.argtypes = ctypes.c_int32, [ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern int mprotect(void *__addr, size_t __len, int __prot) __attribute__((nothrow))
 try: (mprotect:=dll.mprotect).restype, mprotect.argtypes = ctypes.c_int32, [ctypes.c_void_p, size_t, ctypes.c_int32]
 except AttributeError: pass
 
-# extern int msync(void *__addr, size_t __len, int __flags)
 try: (msync:=dll.msync).restype, msync.argtypes = ctypes.c_int32, [ctypes.c_void_p, size_t, ctypes.c_int32]
 except AttributeError: pass
 
-# extern int madvise(void *__addr, size_t __len, int __advice) __attribute__((nothrow))
 try: (madvise:=dll.madvise).restype, madvise.argtypes = ctypes.c_int32, [ctypes.c_void_p, size_t, ctypes.c_int32]
 except AttributeError: pass
 
-# extern int posix_madvise(void *__addr, size_t __len, int __advice) __attribute__((nothrow))
 try: (posix_madvise:=dll.posix_madvise).restype, posix_madvise.argtypes = ctypes.c_int32, [ctypes.c_void_p, size_t, ctypes.c_int32]
 except AttributeError: pass
 
-# extern int mlock(const void *__addr, size_t __len) __attribute__((nothrow))
 try: (mlock:=dll.mlock).restype, mlock.argtypes = ctypes.c_int32, [ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern int munlock(const void *__addr, size_t __len) __attribute__((nothrow))
 try: (munlock:=dll.munlock).restype, munlock.argtypes = ctypes.c_int32, [ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern int mlockall(int __flags) __attribute__((nothrow))
 try: (mlockall:=dll.mlockall).restype, mlockall.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern int munlockall(void) __attribute__((nothrow))
 try: (munlockall:=dll.munlockall).restype, munlockall.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern int mincore(void *__start, size_t __len, unsigned char *__vec) __attribute__((nothrow))
 try: (mincore:=dll.mincore).restype, mincore.argtypes = ctypes.c_int32, [ctypes.c_void_p, size_t, ctypes.POINTER(ctypes.c_ubyte)]
 except AttributeError: pass
 
-# extern int shm_open(const char *__name, int __oflag, mode_t __mode)
 try: (shm_open:=dll.shm_open).restype, shm_open.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.c_int32, mode_t]
 except AttributeError: pass
 
-# extern int shm_unlink(const char *__name)
 try: (shm_unlink:=dll.shm_unlink).restype, shm_unlink.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern void *memcpy(void *restrict __dest, const void *restrict __src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (memcpy:=dll.memcpy).restype, memcpy.argtypes = ctypes.c_void_p, [ctypes.c_void_p, ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern void *memmove(void *__dest, const void *__src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (memmove:=dll.memmove).restype, memmove.argtypes = ctypes.c_void_p, [ctypes.c_void_p, ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern void *memccpy(void *restrict __dest, const void *restrict __src, int __c, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (memccpy:=dll.memccpy).restype, memccpy.argtypes = ctypes.c_void_p, [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, size_t]
 except AttributeError: pass
 
-# extern void *memset(void *__s, int __c, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (memset:=dll.memset).restype, memset.argtypes = ctypes.c_void_p, [ctypes.c_void_p, ctypes.c_int32, size_t]
 except AttributeError: pass
 
-# extern int memcmp(const void *__s1, const void *__s2, size_t __n) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2)))
 try: (memcmp:=dll.memcmp).restype, memcmp.argtypes = ctypes.c_int32, [ctypes.c_void_p, ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern int __memcmpeq(const void *__s1, const void *__s2, size_t __n) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2)))
 try: (__memcmpeq:=dll.__memcmpeq).restype, __memcmpeq.argtypes = ctypes.c_int32, [ctypes.c_void_p, ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern void *memchr(const void *__s, int __c, size_t __n) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1)))
 try: (memchr:=dll.memchr).restype, memchr.argtypes = ctypes.c_void_p, [ctypes.c_void_p, ctypes.c_int32, size_t]
 except AttributeError: pass
 
-# extern char *strcpy(char *restrict __dest, const char *restrict __src) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (strcpy:=dll.strcpy).restype, strcpy.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *strncpy(char *restrict __dest, const char *restrict __src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (strncpy:=dll.strncpy).restype, strncpy.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern char *strcat(char *restrict __dest, const char *restrict __src) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (strcat:=dll.strcat).restype, strcat.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *strncat(char *restrict __dest, const char *restrict __src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (strncat:=dll.strncat).restype, strncat.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern int strcmp(const char *__s1, const char *__s2) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2)))
 try: (strcmp:=dll.strcmp).restype, strcmp.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int strncmp(const char *__s1, const char *__s2, size_t __n) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2)))
 try: (strncmp:=dll.strncmp).restype, strncmp.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern int strcoll(const char *__s1, const char *__s2) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2)))
 try: (strcoll:=dll.strcoll).restype, strcoll.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern unsigned long strxfrm(char *restrict __dest, const char *restrict __src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(2)))
 try: (strxfrm:=dll.strxfrm).restype, strxfrm.argtypes = ctypes.c_uint64, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
@@ -135,131 +107,99 @@ struct___locale_struct._fields_ = [
   ('__names', (ctypes.POINTER(ctypes.c_char) * 13)),
 ]
 locale_t = ctypes.POINTER(struct___locale_struct)
-# extern int strcoll_l(const char *__s1, const char *__s2, locale_t __l) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2, 3)))
 try: (strcoll_l:=dll.strcoll_l).restype, strcoll_l.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), locale_t]
 except AttributeError: pass
 
-# extern size_t strxfrm_l(char *__dest, const char *__src, size_t __n, locale_t __l) __attribute__((nothrow)) __attribute__((nonnull(2, 4)))
 try: (strxfrm_l:=dll.strxfrm_l).restype, strxfrm_l.argtypes = size_t, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t, locale_t]
 except AttributeError: pass
 
-# extern char *strdup(const char *__s) __attribute__((nothrow)) __attribute__((malloc)) __attribute__((nonnull(1)))
 try: (strdup:=dll.strdup).restype, strdup.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *strndup(const char *__string, size_t __n) __attribute__((nothrow)) __attribute__((malloc)) __attribute__((nonnull(1)))
 try: (strndup:=dll.strndup).restype, strndup.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern char *strchr(const char *__s, int __c) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1)))
 try: (strchr:=dll.strchr).restype, strchr.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.c_int32]
 except AttributeError: pass
 
-# extern char *strrchr(const char *__s, int __c) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1)))
 try: (strrchr:=dll.strrchr).restype, strrchr.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.c_int32]
 except AttributeError: pass
 
-# extern char *strchrnul(const char *__s, int __c) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1)))
 try: (strchrnul:=dll.strchrnul).restype, strchrnul.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.c_int32]
 except AttributeError: pass
 
-# extern unsigned long strcspn(const char *__s, const char *__reject) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2)))
 try: (strcspn:=dll.strcspn).restype, strcspn.argtypes = ctypes.c_uint64, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern unsigned long strspn(const char *__s, const char *__accept) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2)))
 try: (strspn:=dll.strspn).restype, strspn.argtypes = ctypes.c_uint64, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *strpbrk(const char *__s, const char *__accept) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2)))
 try: (strpbrk:=dll.strpbrk).restype, strpbrk.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *strstr(const char *__haystack, const char *__needle) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2)))
 try: (strstr:=dll.strstr).restype, strstr.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *strtok(char *restrict __s, const char *restrict __delim) __attribute__((nothrow)) __attribute__((nonnull(2)))
 try: (strtok:=dll.strtok).restype, strtok.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *__strtok_r(char *restrict __s, const char *restrict __delim, char **restrict __save_ptr) __attribute__((nothrow)) __attribute__((nonnull(2, 3)))
 try: (__strtok_r:=dll.__strtok_r).restype, __strtok_r.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
 except AttributeError: pass
 
-# extern char *strtok_r(char *restrict __s, const char *restrict __delim, char **restrict __save_ptr) __attribute__((nothrow)) __attribute__((nonnull(2, 3)))
 try: (strtok_r:=dll.strtok_r).restype, strtok_r.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
 except AttributeError: pass
 
-# extern char *strcasestr(const char *__haystack, const char *__needle) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 2)))
 try: (strcasestr:=dll.strcasestr).restype, strcasestr.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern void *memmem(const void *__haystack, size_t __haystacklen, const void *__needle, size_t __needlelen) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1, 3)))
 try: (memmem:=dll.memmem).restype, memmem.argtypes = ctypes.c_void_p, [ctypes.c_void_p, size_t, ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern void *__mempcpy(void *restrict __dest, const void *restrict __src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (__mempcpy:=dll.__mempcpy).restype, __mempcpy.argtypes = ctypes.c_void_p, [ctypes.c_void_p, ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern void *mempcpy(void *restrict __dest, const void *restrict __src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (mempcpy:=dll.mempcpy).restype, mempcpy.argtypes = ctypes.c_void_p, [ctypes.c_void_p, ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern unsigned long strlen(const char *__s) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1)))
 try: (strlen:=dll.strlen).restype, strlen.argtypes = ctypes.c_uint64, [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern size_t strnlen(const char *__string, size_t __maxlen) __attribute__((nothrow)) __attribute__((pure)) __attribute__((nonnull(1)))
 try: (strnlen:=dll.strnlen).restype, strnlen.argtypes = size_t, [ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern char *strerror(int __errnum) __attribute__((nothrow))
 try: (strerror:=dll.strerror).restype, strerror.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.c_int32]
 except AttributeError: pass
 
-# extern int strerror_r(int __errnum, char *__buf, size_t __buflen) asm("__xpg_strerror_r") __attribute__((nothrow)) __attribute__((nonnull(2)))
 try: (strerror_r:=dll.strerror_r).restype, strerror_r.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern char *strerror_l(int __errnum, locale_t __l) __attribute__((nothrow))
 try: (strerror_l:=dll.strerror_l).restype, strerror_l.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.c_int32, locale_t]
 except AttributeError: pass
 
-# extern void explicit_bzero(void *__s, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (explicit_bzero:=dll.explicit_bzero).restype, explicit_bzero.argtypes = None, [ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern char *strsep(char **restrict __stringp, const char *restrict __delim) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (strsep:=dll.strsep).restype, strsep.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.POINTER(ctypes.c_char)), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *strsignal(int __sig) __attribute__((nothrow))
 try: (strsignal:=dll.strsignal).restype, strsignal.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.c_int32]
 except AttributeError: pass
 
-# extern char *__stpcpy(char *restrict __dest, const char *restrict __src) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (__stpcpy:=dll.__stpcpy).restype, __stpcpy.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *stpcpy(char *restrict __dest, const char *restrict __src) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (stpcpy:=dll.stpcpy).restype, stpcpy.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *__stpncpy(char *restrict __dest, const char *restrict __src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (__stpncpy:=dll.__stpncpy).restype, __stpncpy.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern char *stpncpy(char *restrict __dest, const char *restrict __src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (stpncpy:=dll.stpncpy).restype, stpncpy.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern size_t strlcpy(char *restrict __dest, const char *restrict __src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (strlcpy:=dll.strlcpy).restype, strlcpy.argtypes = size_t, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern size_t strlcat(char *restrict __dest, const char *restrict __src, size_t __n) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (strlcat:=dll.strlcat).restype, strlcat.argtypes = size_t, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
@@ -640,437 +580,330 @@ useconds_t = ctypes.c_uint32
 pid_t = ctypes.c_int32
 intptr_t = ctypes.c_int64
 socklen_t = ctypes.c_uint32
-# extern int access(const char *__name, int __type) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (access:=dll.access).restype, access.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.c_int32]
 except AttributeError: pass
 
-# extern int faccessat(int __fd, const char *__file, int __type, int __flag) __attribute__((nothrow)) __attribute__((nonnull(2)))
 try: (faccessat:=dll.faccessat).restype, faccessat.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.POINTER(ctypes.c_char), ctypes.c_int32, ctypes.c_int32]
 except AttributeError: pass
 
-# extern __off_t lseek(int __fd, __off_t __offset, int __whence) __attribute__((nothrow))
 try: (lseek:=dll.lseek).restype, lseek.argtypes = ctypes.c_int64, [ctypes.c_int32, ctypes.c_int64, ctypes.c_int32]
 except AttributeError: pass
 
-# extern int close(int __fd)
 try: (close:=dll.close).restype, close.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern void closefrom(int __lowfd) __attribute__((nothrow))
 try: (closefrom:=dll.closefrom).restype, closefrom.argtypes = None, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern ssize_t read(int __fd, void *__buf, size_t __nbytes)
 try: (read:=dll.read).restype, read.argtypes = ssize_t, [ctypes.c_int32, ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern ssize_t write(int __fd, const void *__buf, size_t __n)
 try: (write:=dll.write).restype, write.argtypes = ssize_t, [ctypes.c_int32, ctypes.c_void_p, size_t]
 except AttributeError: pass
 
-# extern ssize_t pread(int __fd, void *__buf, size_t __nbytes, __off_t __offset)
 try: (pread:=dll.pread).restype, pread.argtypes = ssize_t, [ctypes.c_int32, ctypes.c_void_p, size_t, ctypes.c_int64]
 except AttributeError: pass
 
-# extern ssize_t pwrite(int __fd, const void *__buf, size_t __n, __off_t __offset)
 try: (pwrite:=dll.pwrite).restype, pwrite.argtypes = ssize_t, [ctypes.c_int32, ctypes.c_void_p, size_t, ctypes.c_int64]
 except AttributeError: pass
 
-# extern int pipe(int __pipedes[2]) __attribute__((nothrow))
 try: (pipe:=dll.pipe).restype, pipe.argtypes = ctypes.c_int32, [(ctypes.c_int32 * 2)]
 except AttributeError: pass
 
-# extern unsigned int alarm(unsigned int __seconds) __attribute__((nothrow))
 try: (alarm:=dll.alarm).restype, alarm.argtypes = ctypes.c_uint32, [ctypes.c_uint32]
 except AttributeError: pass
 
-# extern unsigned int sleep(unsigned int __seconds)
 try: (sleep:=dll.sleep).restype, sleep.argtypes = ctypes.c_uint32, [ctypes.c_uint32]
 except AttributeError: pass
 
 __useconds_t = ctypes.c_uint32
-# extern __useconds_t ualarm(__useconds_t __value, __useconds_t __interval) __attribute__((nothrow))
 try: (ualarm:=dll.ualarm).restype, ualarm.argtypes = ctypes.c_uint32, [ctypes.c_uint32, ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int usleep(__useconds_t __useconds)
 try: (usleep:=dll.usleep).restype, usleep.argtypes = ctypes.c_int32, [ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int pause(void)
 try: (pause:=dll.pause).restype, pause.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
 __uid_t = ctypes.c_uint32
 __gid_t = ctypes.c_uint32
-# extern int chown(const char *__file, __uid_t __owner, __gid_t __group) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (chown:=dll.chown).restype, chown.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.c_uint32, ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int fchown(int __fd, __uid_t __owner, __gid_t __group) __attribute__((nothrow))
 try: (fchown:=dll.fchown).restype, fchown.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.c_uint32, ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int lchown(const char *__file, __uid_t __owner, __gid_t __group) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (lchown:=dll.lchown).restype, lchown.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.c_uint32, ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int fchownat(int __fd, const char *__file, __uid_t __owner, __gid_t __group, int __flag) __attribute__((nothrow)) __attribute__((nonnull(2)))
 try: (fchownat:=dll.fchownat).restype, fchownat.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.POINTER(ctypes.c_char), ctypes.c_uint32, ctypes.c_uint32, ctypes.c_int32]
 except AttributeError: pass
 
-# extern int chdir(const char *__path) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (chdir:=dll.chdir).restype, chdir.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int fchdir(int __fd) __attribute__((nothrow))
 try: (fchdir:=dll.fchdir).restype, fchdir.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern char *getcwd(char *__buf, size_t __size) __attribute__((nothrow))
 try: (getcwd:=dll.getcwd).restype, getcwd.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern char *getwd(char *__buf) __attribute__((nothrow)) __attribute__((nonnull(1))) __attribute__((deprecated("")))
 try: (getwd:=dll.getwd).restype, getwd.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int dup(int __fd) __attribute__((nothrow))
 try: (dup:=dll.dup).restype, dup.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern int dup2(int __fd, int __fd2) __attribute__((nothrow))
 try: (dup2:=dll.dup2).restype, dup2.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.c_int32]
 except AttributeError: pass
 
 try: __environ = ctypes.POINTER(ctypes.POINTER(ctypes.c_char)).in_dll(dll, '__environ')
 except (ValueError,AttributeError): pass
-# extern int execve(const char *__path, char *const __argv[], char *const __envp[]) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (execve:=dll.execve).restype, execve.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), (ctypes.POINTER(ctypes.c_char) * 0), (ctypes.POINTER(ctypes.c_char) * 0)]
 except AttributeError: pass
 
-# extern int fexecve(int __fd, char *const __argv[], char *const __envp[]) __attribute__((nothrow)) __attribute__((nonnull(2)))
 try: (fexecve:=dll.fexecve).restype, fexecve.argtypes = ctypes.c_int32, [ctypes.c_int32, (ctypes.POINTER(ctypes.c_char) * 0), (ctypes.POINTER(ctypes.c_char) * 0)]
 except AttributeError: pass
 
-# extern int execv(const char *__path, char *const __argv[]) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (execv:=dll.execv).restype, execv.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), (ctypes.POINTER(ctypes.c_char) * 0)]
 except AttributeError: pass
 
-# extern int execle(const char *__path, const char *__arg, ...) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (execle:=dll.execle).restype, execle.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int execl(const char *__path, const char *__arg, ...) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (execl:=dll.execl).restype, execl.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int execvp(const char *__file, char *const __argv[]) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (execvp:=dll.execvp).restype, execvp.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), (ctypes.POINTER(ctypes.c_char) * 0)]
 except AttributeError: pass
 
-# extern int execlp(const char *__file, const char *__arg, ...) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (execlp:=dll.execlp).restype, execlp.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int nice(int __inc) __attribute__((nothrow))
 try: (nice:=dll.nice).restype, nice.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern void _exit(int __status)
 try: (_exit:=dll._exit).restype, _exit.argtypes = None, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern long pathconf(const char *__path, int __name) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (pathconf:=dll.pathconf).restype, pathconf.argtypes = ctypes.c_int64, [ctypes.POINTER(ctypes.c_char), ctypes.c_int32]
 except AttributeError: pass
 
-# extern long fpathconf(int __fd, int __name) __attribute__((nothrow))
 try: (fpathconf:=dll.fpathconf).restype, fpathconf.argtypes = ctypes.c_int64, [ctypes.c_int32, ctypes.c_int32]
 except AttributeError: pass
 
-# extern long sysconf(int __name) __attribute__((nothrow))
 try: (sysconf:=dll.sysconf).restype, sysconf.argtypes = ctypes.c_int64, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern size_t confstr(int __name, char *__buf, size_t __len) __attribute__((nothrow))
 try: (confstr:=dll.confstr).restype, confstr.argtypes = size_t, [ctypes.c_int32, ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
 __pid_t = ctypes.c_int32
-# extern __pid_t getpid(void) __attribute__((nothrow))
 try: (getpid:=dll.getpid).restype, getpid.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern __pid_t getppid(void) __attribute__((nothrow))
 try: (getppid:=dll.getppid).restype, getppid.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern __pid_t getpgrp(void) __attribute__((nothrow))
 try: (getpgrp:=dll.getpgrp).restype, getpgrp.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern __pid_t __getpgid(__pid_t __pid) __attribute__((nothrow))
 try: (__getpgid:=dll.__getpgid).restype, __getpgid.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern __pid_t getpgid(__pid_t __pid) __attribute__((nothrow))
 try: (getpgid:=dll.getpgid).restype, getpgid.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern int setpgid(__pid_t __pid, __pid_t __pgid) __attribute__((nothrow))
 try: (setpgid:=dll.setpgid).restype, setpgid.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.c_int32]
 except AttributeError: pass
 
-# extern int setpgrp(void) __attribute__((nothrow))
 try: (setpgrp:=dll.setpgrp).restype, setpgrp.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern __pid_t setsid(void) __attribute__((nothrow))
 try: (setsid:=dll.setsid).restype, setsid.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern __pid_t getsid(__pid_t __pid) __attribute__((nothrow))
 try: (getsid:=dll.getsid).restype, getsid.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern __uid_t getuid(void) __attribute__((nothrow))
 try: (getuid:=dll.getuid).restype, getuid.argtypes = ctypes.c_uint32, []
 except AttributeError: pass
 
-# extern __uid_t geteuid(void) __attribute__((nothrow))
 try: (geteuid:=dll.geteuid).restype, geteuid.argtypes = ctypes.c_uint32, []
 except AttributeError: pass
 
-# extern __gid_t getgid(void) __attribute__((nothrow))
 try: (getgid:=dll.getgid).restype, getgid.argtypes = ctypes.c_uint32, []
 except AttributeError: pass
 
-# extern __gid_t getegid(void) __attribute__((nothrow))
 try: (getegid:=dll.getegid).restype, getegid.argtypes = ctypes.c_uint32, []
 except AttributeError: pass
 
-# extern int getgroups(int __size, __gid_t __list[]) __attribute__((nothrow))
 try: (getgroups:=dll.getgroups).restype, getgroups.argtypes = ctypes.c_int32, [ctypes.c_int32, (ctypes.c_uint32 * 0)]
 except AttributeError: pass
 
-# extern int setuid(__uid_t __uid) __attribute__((nothrow))
 try: (setuid:=dll.setuid).restype, setuid.argtypes = ctypes.c_int32, [ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int setreuid(__uid_t __ruid, __uid_t __euid) __attribute__((nothrow))
 try: (setreuid:=dll.setreuid).restype, setreuid.argtypes = ctypes.c_int32, [ctypes.c_uint32, ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int seteuid(__uid_t __uid) __attribute__((nothrow))
 try: (seteuid:=dll.seteuid).restype, seteuid.argtypes = ctypes.c_int32, [ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int setgid(__gid_t __gid) __attribute__((nothrow))
 try: (setgid:=dll.setgid).restype, setgid.argtypes = ctypes.c_int32, [ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int setregid(__gid_t __rgid, __gid_t __egid) __attribute__((nothrow))
 try: (setregid:=dll.setregid).restype, setregid.argtypes = ctypes.c_int32, [ctypes.c_uint32, ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int setegid(__gid_t __gid) __attribute__((nothrow))
 try: (setegid:=dll.setegid).restype, setegid.argtypes = ctypes.c_int32, [ctypes.c_uint32]
 except AttributeError: pass
 
-# extern __pid_t fork(void) __attribute__((nothrow))
 try: (fork:=dll.fork).restype, fork.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern int vfork(void) __attribute__((nothrow))
 try: (vfork:=dll.vfork).restype, vfork.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern char *ttyname(int __fd) __attribute__((nothrow))
 try: (ttyname:=dll.ttyname).restype, ttyname.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.c_int32]
 except AttributeError: pass
 
-# extern int ttyname_r(int __fd, char *__buf, size_t __buflen) __attribute__((nothrow)) __attribute__((nonnull(2)))
 try: (ttyname_r:=dll.ttyname_r).restype, ttyname_r.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern int isatty(int __fd) __attribute__((nothrow))
 try: (isatty:=dll.isatty).restype, isatty.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern int ttyslot(void) __attribute__((nothrow))
 try: (ttyslot:=dll.ttyslot).restype, ttyslot.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern int link(const char *__from, const char *__to) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (link:=dll.link).restype, link.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int linkat(int __fromfd, const char *__from, int __tofd, const char *__to, int __flags) __attribute__((nothrow)) __attribute__((nonnull(2, 4)))
 try: (linkat:=dll.linkat).restype, linkat.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.POINTER(ctypes.c_char), ctypes.c_int32, ctypes.POINTER(ctypes.c_char), ctypes.c_int32]
 except AttributeError: pass
 
-# extern int symlink(const char *__from, const char *__to) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (symlink:=dll.symlink).restype, symlink.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern ssize_t readlink(const char *restrict __path, char *restrict __buf, size_t __len) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (readlink:=dll.readlink).restype, readlink.argtypes = ssize_t, [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern int symlinkat(const char *__from, int __tofd, const char *__to) __attribute__((nothrow)) __attribute__((nonnull(1, 3)))
 try: (symlinkat:=dll.symlinkat).restype, symlinkat.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.c_int32, ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern ssize_t readlinkat(int __fd, const char *restrict __path, char *restrict __buf, size_t __len) __attribute__((nothrow)) __attribute__((nonnull(2, 3)))
 try: (readlinkat:=dll.readlinkat).restype, readlinkat.argtypes = ssize_t, [ctypes.c_int32, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern int unlink(const char *__name) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (unlink:=dll.unlink).restype, unlink.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int unlinkat(int __fd, const char *__name, int __flag) __attribute__((nothrow)) __attribute__((nonnull(2)))
 try: (unlinkat:=dll.unlinkat).restype, unlinkat.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.POINTER(ctypes.c_char), ctypes.c_int32]
 except AttributeError: pass
 
-# extern int rmdir(const char *__path) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (rmdir:=dll.rmdir).restype, rmdir.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern __pid_t tcgetpgrp(int __fd) __attribute__((nothrow))
 try: (tcgetpgrp:=dll.tcgetpgrp).restype, tcgetpgrp.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern int tcsetpgrp(int __fd, __pid_t __pgrp_id) __attribute__((nothrow))
 try: (tcsetpgrp:=dll.tcsetpgrp).restype, tcsetpgrp.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.c_int32]
 except AttributeError: pass
 
-# extern char *getlogin(void)
 try: (getlogin:=dll.getlogin).restype, getlogin.argtypes = ctypes.POINTER(ctypes.c_char), []
 except AttributeError: pass
 
-# extern int getlogin_r(char *__name, size_t __name_len) __attribute__((nonnull(1)))
 try: (getlogin_r:=dll.getlogin_r).restype, getlogin_r.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern int setlogin(const char *__name) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (setlogin:=dll.setlogin).restype, setlogin.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int gethostname(char *__name, size_t __len) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (gethostname:=dll.gethostname).restype, gethostname.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern int sethostname(const char *__name, size_t __len) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (sethostname:=dll.sethostname).restype, sethostname.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern int sethostid(long __id) __attribute__((nothrow))
 try: (sethostid:=dll.sethostid).restype, sethostid.argtypes = ctypes.c_int32, [ctypes.c_int64]
 except AttributeError: pass
 
-# extern int getdomainname(char *__name, size_t __len) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (getdomainname:=dll.getdomainname).restype, getdomainname.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern int setdomainname(const char *__name, size_t __len) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (setdomainname:=dll.setdomainname).restype, setdomainname.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), size_t]
 except AttributeError: pass
 
-# extern int vhangup(void) __attribute__((nothrow))
 try: (vhangup:=dll.vhangup).restype, vhangup.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern int revoke(const char *__file) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (revoke:=dll.revoke).restype, revoke.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int profil(unsigned short *__sample_buffer, size_t __size, size_t __offset, unsigned int __scale) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (profil:=dll.profil).restype, profil.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_uint16), size_t, size_t, ctypes.c_uint32]
 except AttributeError: pass
 
-# extern int acct(const char *__name) __attribute__((nothrow))
 try: (acct:=dll.acct).restype, acct.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *getusershell(void) __attribute__((nothrow))
 try: (getusershell:=dll.getusershell).restype, getusershell.argtypes = ctypes.POINTER(ctypes.c_char), []
 except AttributeError: pass
 
-# extern void endusershell(void) __attribute__((nothrow))
 try: (endusershell:=dll.endusershell).restype, endusershell.argtypes = None, []
 except AttributeError: pass
 
-# extern void setusershell(void) __attribute__((nothrow))
 try: (setusershell:=dll.setusershell).restype, setusershell.argtypes = None, []
 except AttributeError: pass
 
-# extern int daemon(int __nochdir, int __noclose) __attribute__((nothrow))
 try: (daemon:=dll.daemon).restype, daemon.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.c_int32]
 except AttributeError: pass
 
-# extern int chroot(const char *__path) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (chroot:=dll.chroot).restype, chroot.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern char *getpass(const char *__prompt) __attribute__((nonnull(1)))
 try: (getpass:=dll.getpass).restype, getpass.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# extern int fsync(int __fd)
 try: (fsync:=dll.fsync).restype, fsync.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern long gethostid(void)
 try: (gethostid:=dll.gethostid).restype, gethostid.argtypes = ctypes.c_int64, []
 except AttributeError: pass
 
-# extern void sync(void) __attribute__((nothrow))
 try: (sync:=dll.sync).restype, sync.argtypes = None, []
 except AttributeError: pass
 
-# extern int getpagesize(void) __attribute__((nothrow)) __attribute__((const))
 try: (getpagesize:=dll.getpagesize).restype, getpagesize.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern int getdtablesize(void) __attribute__((nothrow))
 try: (getdtablesize:=dll.getdtablesize).restype, getdtablesize.argtypes = ctypes.c_int32, []
 except AttributeError: pass
 
-# extern int truncate(const char *__file, __off_t __length) __attribute__((nothrow)) __attribute__((nonnull(1)))
 try: (truncate:=dll.truncate).restype, truncate.argtypes = ctypes.c_int32, [ctypes.POINTER(ctypes.c_char), ctypes.c_int64]
 except AttributeError: pass
 
-# extern int ftruncate(int __fd, __off_t __length) __attribute__((nothrow))
 try: (ftruncate:=dll.ftruncate).restype, ftruncate.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.c_int64]
 except AttributeError: pass
 
-# extern int brk(void *__addr) __attribute__((nothrow))
 try: (brk:=dll.brk).restype, brk.argtypes = ctypes.c_int32, [ctypes.c_void_p]
 except AttributeError: pass
 
-# extern void *sbrk(intptr_t __delta) __attribute__((nothrow))
 try: (sbrk:=dll.sbrk).restype, sbrk.argtypes = ctypes.c_void_p, [intptr_t]
 except AttributeError: pass
 
-# extern long syscall(long __sysno, ...) __attribute__((nothrow))
 try: (syscall:=dll.syscall).restype, syscall.argtypes = ctypes.c_int64, [ctypes.c_int64]
 except AttributeError: pass
 
-# extern int lockf(int __fd, int __cmd, __off_t __len)
 try: (lockf:=dll.lockf).restype, lockf.argtypes = ctypes.c_int32, [ctypes.c_int32, ctypes.c_int32, ctypes.c_int64]
 except AttributeError: pass
 
-# extern int fdatasync(int __fildes)
 try: (fdatasync:=dll.fdatasync).restype, fdatasync.argtypes = ctypes.c_int32, [ctypes.c_int32]
 except AttributeError: pass
 
-# extern char *crypt(const char *__key, const char *__salt) __attribute__((nothrow)) __attribute__((nonnull(1, 2)))
 try: (crypt:=dll.crypt).restype, crypt.argtypes = ctypes.POINTER(ctypes.c_char), [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)]
 except AttributeError: pass
 
-# int getentropy(void *__buffer, size_t __length)
 try: (getentropy:=dll.getentropy).restype, getentropy.argtypes = ctypes.c_int32, [ctypes.c_void_p, size_t]
 except AttributeError: pass
 

tinygrad/runtime/autogen/nvjitlink.py

@@ -0,0 +1,75 @@
+# 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
+try: (nvJitLinkCreate:=dll.nvJitLinkCreate).restype, nvJitLinkCreate.argtypes = nvJitLinkResult, [ctypes.POINTER(nvJitLinkHandle), uint32_t, ctypes.POINTER(ctypes.POINTER(ctypes.c_char))]
+except AttributeError: pass
+
+try: (nvJitLinkDestroy:=dll.nvJitLinkDestroy).restype, nvJitLinkDestroy.argtypes = nvJitLinkResult, [ctypes.POINTER(nvJitLinkHandle)]
+except AttributeError: pass
+
+size_t = ctypes.c_uint64
+try: (nvJitLinkAddData:=dll.nvJitLinkAddData).restype, nvJitLinkAddData.argtypes = nvJitLinkResult, [nvJitLinkHandle, nvJitLinkInputType, ctypes.c_void_p, size_t, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+try: (nvJitLinkAddFile:=dll.nvJitLinkAddFile).restype, nvJitLinkAddFile.argtypes = nvJitLinkResult, [nvJitLinkHandle, nvJitLinkInputType, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+try: (nvJitLinkComplete:=dll.nvJitLinkComplete).restype, nvJitLinkComplete.argtypes = nvJitLinkResult, [nvJitLinkHandle]
+except AttributeError: pass
+
+try: (nvJitLinkGetLinkedCubinSize:=dll.nvJitLinkGetLinkedCubinSize).restype, nvJitLinkGetLinkedCubinSize.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(size_t)]
+except AttributeError: pass
+
+try: (nvJitLinkGetLinkedCubin:=dll.nvJitLinkGetLinkedCubin).restype, nvJitLinkGetLinkedCubin.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.c_void_p]
+except AttributeError: pass
+
+try: (nvJitLinkGetLinkedPtxSize:=dll.nvJitLinkGetLinkedPtxSize).restype, nvJitLinkGetLinkedPtxSize.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(size_t)]
+except AttributeError: pass
+
+try: (nvJitLinkGetLinkedPtx:=dll.nvJitLinkGetLinkedPtx).restype, nvJitLinkGetLinkedPtx.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+try: (nvJitLinkGetErrorLogSize:=dll.nvJitLinkGetErrorLogSize).restype, nvJitLinkGetErrorLogSize.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(size_t)]
+except AttributeError: pass
+
+try: (nvJitLinkGetErrorLog:=dll.nvJitLinkGetErrorLog).restype, nvJitLinkGetErrorLog.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+try: (nvJitLinkGetInfoLogSize:=dll.nvJitLinkGetInfoLogSize).restype, nvJitLinkGetInfoLogSize.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(size_t)]
+except AttributeError: pass
+
+try: (nvJitLinkGetInfoLog:=dll.nvJitLinkGetInfoLog).restype, nvJitLinkGetInfoLog.argtypes = nvJitLinkResult, [nvJitLinkHandle, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+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,92 @@
 # 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
+try: (nvrtcGetErrorString:=dll.nvrtcGetErrorString).restype, nvrtcGetErrorString.argtypes = ctypes.POINTER(ctypes.c_char), [nvrtcResult]
+except AttributeError: pass
 
+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))
+try: (nvrtcGetNumSupportedArchs:=dll.nvrtcGetNumSupportedArchs).restype, nvrtcGetNumSupportedArchs.argtypes = nvrtcResult, [ctypes.POINTER(ctypes.c_int32)]
+except AttributeError: pass
 
+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
+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
+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
+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']
+try: (nvrtcGetPTXSize:=dll.nvrtcGetPTXSize).restype, nvrtcGetPTXSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
+except AttributeError: pass
+
+try: (nvrtcGetPTX:=dll.nvrtcGetPTX).restype, nvrtcGetPTX.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+try: (nvrtcGetCUBINSize:=dll.nvrtcGetCUBINSize).restype, nvrtcGetCUBINSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
+except AttributeError: pass
+
+try: (nvrtcGetCUBIN:=dll.nvrtcGetCUBIN).restype, nvrtcGetCUBIN.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+try: (nvrtcGetNVVMSize:=dll.nvrtcGetNVVMSize).restype, nvrtcGetNVVMSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
+except AttributeError: pass
+
+try: (nvrtcGetNVVM:=dll.nvrtcGetNVVM).restype, nvrtcGetNVVM.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+try: (nvrtcGetLTOIRSize:=dll.nvrtcGetLTOIRSize).restype, nvrtcGetLTOIRSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
+except AttributeError: pass
+
+try: (nvrtcGetLTOIR:=dll.nvrtcGetLTOIR).restype, nvrtcGetLTOIR.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+try: (nvrtcGetOptiXIRSize:=dll.nvrtcGetOptiXIRSize).restype, nvrtcGetOptiXIRSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
+except AttributeError: pass
+
+try: (nvrtcGetOptiXIR:=dll.nvrtcGetOptiXIR).restype, nvrtcGetOptiXIR.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+try: (nvrtcGetProgramLogSize:=dll.nvrtcGetProgramLogSize).restype, nvrtcGetProgramLogSize.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(size_t)]
+except AttributeError: pass
+
+try: (nvrtcGetProgramLog:=dll.nvrtcGetProgramLog).restype, nvrtcGetProgramLog.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+try: (nvrtcAddNameExpression:=dll.nvrtcAddNameExpression).restype, nvrtcAddNameExpression.argtypes = nvrtcResult, [nvrtcProgram, ctypes.POINTER(ctypes.c_char)]
+except AttributeError: pass
+
+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,256 @@
-# 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)
+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
+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
+
+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)
+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/graph/metal.py

@@ -5,15 +5,9 @@ from tinygrad.helpers import dedup, getenv, merge_dicts, PROFILE
 from tinygrad.device import Buffer, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.engine.realize import ExecItem, CompiledRunner
 from tinygrad.engine.jit import GraphRunner, GraphException
-from tinygrad.runtime.ops_metal import wait_check, msg, libobjc, to_struct, objc_instance,\
-  MTLResourceOptions, cmdbuf_st_time, cmdbuf_en_time, objc_id, to_ns_str
-
-class MTLIndirectCommandType:
-  MTLIndirectCommandTypeConcurrentDispatch = (1 << 5)
-
-class MTLResourceUsage:
-  MTLResourceUsageRead = 0b01
-  MTLResourceUsageWrite = 0b10
+from tinygrad.runtime.ops_metal import wait_check, to_ns_str
+from tinygrad.runtime.autogen import metal
+from tinygrad.runtime.support import objc
 
 class MetalGraph(GraphRunner):
   def __init__(self, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[str, int]):
@@ -21,16 +15,17 @@ class MetalGraph(GraphRunner):
     if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException
 
     # create metal batch exec
-    icb_descriptor = msg("new", objc_instance)(libobjc.objc_getClass(b"MTLIndirectCommandBufferDescriptor"))
-    msg("setCommandTypes:")(icb_descriptor, MTLIndirectCommandType.MTLIndirectCommandTypeConcurrentDispatch)
-    msg("setInheritBuffers:")(icb_descriptor, False)
-    msg("setInheritPipelineState:")(icb_descriptor, False)
-    msg("setMaxKernelBufferBindCount:")(icb_descriptor, 31)
+    icb_descriptor = metal.MTLIndirectCommandBufferDescriptor.new()
+    icb_descriptor.setCommandTypes(metal.MTLIndirectCommandTypeConcurrentDispatch)
+    icb_descriptor.setInheritBuffers(False)
+    icb_descriptor.setInheritPipelineState(False)
+    icb_descriptor.setMaxKernelBufferBindCount(31)
 
-    self.icb = msg("newIndirectCommandBufferWithDescriptor:maxCommandCount:options:", objc_instance)(self.dev.sysdevice,
-      icb_descriptor, len(jit_cache), MTLResourceOptions.MTLResourceCPUCacheModeDefaultCache)
+    self.icb = self.dev.sysdevice.newIndirectCommandBufferWithDescriptor_maxCommandCount_options(icb_descriptor, len(jit_cache),
+                                                                                                 metal.MTLResourceCPUCacheModeDefaultCache)
     if self.icb.value is None: raise GraphException("create indirect command buffer failed, does your system support this?")
-    icb_label = bytes(msg("UTF8String", ctypes.c_char_p)(msg("description", objc_instance)(self.icb))).decode()
+    # TODO: needs categories
+    icb_label = bytes(objc.msg("UTF8String", ctypes.c_char_p)(objc.msg("description")(self.icb).retained())).decode()
     self.needs_icb_fix = int((m := re.search(r'AGXG(\d+)XFamily', icb_label)) is None or int(m.group(1)) < 15) # not required on M3+
 
     self.fixedvars = merge_dicts([ji.fixedvars for ji in jit_cache])
@@ -40,26 +35,25 @@ class MetalGraph(GraphRunner):
     all_pipelines, all_resources = [], [self.int_buf.buf] if len(self.varlist) else []
     for j,ji in enumerate(jit_cache):
       prg: CompiledRunner = cast(CompiledRunner, ji.prg)
-      icb_command = msg("indirectComputeCommandAtIndex:", objc_instance)(self.icb, j)
+      icb_command = self.icb.indirectComputeCommandAtIndex(j).retained()
       all_pipelines.append(prg._prg.pipeline_state)
-      msg("setComputePipelineState:")(icb_command, prg._prg.pipeline_state)
+      icb_command.setComputePipelineState(prg._prg.pipeline_state)
       for i,b in enumerate(ji.bufs):
         if b is not None and b not in input_rawbuffers:
-          msg("setKernelBuffer:offset:atIndex:")(icb_command, b._buf.buf, b._buf.offset, i)
+          icb_command.setKernelBuffer_offset_atIndex(b._buf.buf, b._buf.offset, i)
           all_resources.append(b._buf.buf)
-      for i,v in enumerate(prg.p.vars):
-        msg("setKernelBuffer:offset:atIndex:")(icb_command, self.int_buf.buf, self.varlist.index(v.expr)*4, len(ji.bufs)+i)
+      for i,v in enumerate(prg.p.vars): icb_command.setKernelBuffer_offset_atIndex(self.int_buf.buf, self.varlist.index(v.expr)*4, len(ji.bufs)+i)
 
       global_size, local_size = prg.p.launch_dims(var_vals)
-      msg("concurrentDispatchThreadgroups:threadsPerThreadgroup:")(icb_command, to_struct(*global_size), to_struct(*local_size))
-      msg("setBarrier")(icb_command)
+      icb_command.concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_size), metal.MTLSize(*local_size))
+      icb_command.setBarrier()
 
     self.all_resources = dedup(all_resources)
     self.all_pipelines = dedup(all_pipelines)
     self.command_buffer: Any = None
     if len(self.varlist): self.int_buf_view = self.dev.allocator._as_buffer(self.int_buf).cast('i')
     for var in self.fixedvars: self.int_buf_view[self.varlist.index(var)] = self.fixedvars[var]
-    self.range = to_struct(0, len(jit_cache))
+    self.range = metal.NSRange(0, len(jit_cache))
 
   def __call__(self, input_rawbuffers: list[Buffer], var_vals: dict[str, int], wait=False) -> float|None:
     if self.command_buffer is not None and self.command_buffer in self.dev.mtl_buffers_in_flight: wait_check(self.command_buffer)
@@ -68,18 +62,18 @@ class MetalGraph(GraphRunner):
 
     all_resources = dedup(self.all_resources + [input_rawbuffers[input_idx]._buf.buf for input_idx in self.input_replace.values()])
     for (j,i),input_idx in self.input_replace.items():
-      computeCommand = msg("indirectComputeCommandAtIndex:", objc_id)(self.icb, j)
-      msg("setKernelBuffer:offset:atIndex:")(computeCommand, input_rawbuffers[input_idx]._buf.buf, input_rawbuffers[input_idx]._buf.offset, i)
+      computeCommand = self.icb.indirectComputeCommandAtIndex(j)
+      computeCommand.setKernelBuffer_offset_atIndex(input_rawbuffers[input_idx]._buf.buf, input_rawbuffers[input_idx]._buf.offset, i)
 
     for j, global_dims, local_dims in self.updated_launch_dims(var_vals):
-      computeCommand = msg("indirectComputeCommandAtIndex:", objc_id)(self.icb, j)
-      msg("concurrentDispatchThreadgroups:threadsPerThreadgroup:")(computeCommand, to_struct(*global_dims), to_struct(*local_dims))
+      computeCommand = self.icb.indirectComputeCommandAtIndex(j)
+      computeCommand.concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_dims), metal.MTLSize(*local_dims))
     for var in self.vars: self.int_buf_view[self.varlist.index(var)] = var_vals[var]
 
-    command_buffer = msg("commandBuffer", objc_instance)(self.dev.mtl_queue)
-    encoder = msg("computeCommandEncoder", objc_instance)(command_buffer)
-    msg("useResources:count:usage:")(encoder, (objc_id * len(all_resources))(*all_resources), len(all_resources),
-        MTLResourceUsage.MTLResourceUsageRead | MTLResourceUsage.MTLResourceUsageWrite)
+    command_buffer = self.dev.mtl_queue.commandBuffer().retained()
+    encoder = command_buffer.computeCommandEncoder().retained()
+    encoder.useResources_count_usage(ctypes.cast((metal.MTLBuffer * len(all_resources))(*all_resources), ctypes.POINTER(metal.MTLResource)),
+                                     len(all_resources), metal.MTLResourceUsageRead | metal.MTLResourceUsageWrite)
 
     # NOTE: the pipelines likely need to be added to the used resources to fix the crash on M1/M2, but I haven't figured out how
     # this is a O(n) hack to get them used. what should work is:
@@ -88,24 +82,24 @@ class MetalGraph(GraphRunner):
     # to repro the crash (which can also crash other running GPU apps), run with FIX_METAL_ICB=0
     if getenv("FIX_METAL_ICB", self.needs_icb_fix):
       for ps in self.all_pipelines:
-        msg("setComputePipelineState:")(encoder, ps)
-        msg("dispatchThreadgroups:threadsPerThreadgroup:")(encoder, to_struct(0,0,0), to_struct(0,0,0))
+        encoder.setComputePipelineState(ps)
+        encoder.dispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(0,0,0), metal.MTLSize(0,0,0))
 
-    msg("executeCommandsInBuffer:withRange:")(encoder, self.icb, self.range)
-    msg("endEncoding")(encoder)
-    msg("setLabel:")(command_buffer, to_ns_str(f"batched {len(self.jit_cache)}"))
-    msg("commit")(command_buffer)
+    encoder.executeCommandsInBuffer_withRange(self.icb, self.range)
+    encoder.endEncoding()
+    command_buffer.setLabel(to_ns_str(f"batched {len(self.jit_cache)}"))
+    command_buffer.commit()
     self.command_buffer = command_buffer
 
     self.dev.mtl_buffers_in_flight.append(command_buffer)
     if wait:
       wait_check(command_buffer)
-      return cmdbuf_en_time(command_buffer) - cmdbuf_st_time(command_buffer)
+      return command_buffer.GPUEndTime() - command_buffer.GPUStartTime()
     return None
 
   def collect_timestamps(self):
     # create a graph event and evenly space each program
-    st, en = decimal.Decimal(cmdbuf_st_time(self.command_buffer)) * 1000000, decimal.Decimal(cmdbuf_en_time(self.command_buffer)) * 1000000
+    st, en = decimal.Decimal(self.command_buffer.GPUStartTime()) * 1000000, decimal.Decimal(self.command_buffer.GPUEndTime()) * 1000000
     ents = [ProfileGraphEntry(self.device, cast(CompiledRunner, ji.prg)._prg.name, i, i+1, is_copy=False) for i,ji in enumerate(self.jit_cache)]
     step = (en-st)/len(ents)
     self.dev.profile_events += [ProfileGraphEvent(ents, [], [st+step*i for i in range(len(ents)+1)])]

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_metal.py

@@ -1,74 +1,42 @@
-import subprocess, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform
-from typing import Any, cast
-from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, T, init_c_struct_t, PROFILE, ProfileRangeEvent, cpu_profile, unwrap
+import subprocess, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform, sys
+from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, init_c_struct_t, PROFILE, ProfileRangeEvent, cpu_profile, unwrap
+import tinygrad.runtime.support.objc as objc
 from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent
 from tinygrad.renderer.cstyle import MetalRenderer
-
-class objc_id(ctypes.c_void_p): # This prevents ctypes from converting response to plain int, and dict.fromkeys() can use it to dedup
-  def __hash__(self): return hash(self.value)
-  def __eq__(self, other): return self.value == other.value
-
-class objc_instance(objc_id): # method with name "new", "alloc" should be freed after use
-  def __del__(self):
-    # CPython doesn't make any guarantees about order in which globals (like `msg` or `libobjc`) are destroyed when the interpreter shuts down
-    # https://github.com/tinygrad/tinygrad/pull/8949 triggered the unlucky ordering which lead to a bunch of errors at exit
-    # TODO: Why isn't `sys.is_finalizing` working?
-    if msg is not None and libobjc is not None: msg("release")(self)
-
-class MTLResourceOptions:
-  MTLResourceCPUCacheModeDefaultCache = 0
-  MTLResourceStorageModeShared = 0 << 4
-
-class MTLPipelineOption:
-  MTLPipelineOptionNone = 0
+from tinygrad.runtime.autogen import metal
 
 # 13 is requestType that metal uses to compile source code into MTLB, there aren't any docs or symbols.
 REQUEST_TYPE_COMPILE = 13
 
-libobjc = ctypes.CDLL("/usr/lib/libobjc.dylib")
-libmetal = ctypes.CDLL("/System/Library/Frameworks/Metal.framework/Metal")
 # Must be loaded for default Metal Device: https://developer.apple.com/documentation/metal/1433401-mtlcreatesystemdefaultdevice?language=objc
 ctypes.CDLL("/System/Library/Frameworks/CoreGraphics.framework/CoreGraphics")
-libdispatch = ctypes.CDLL("/usr/lib/libSystem.dylib") # libdispatch is part of libSystem on mac
-libobjc.objc_getClass.restype = objc_id
-libobjc.sel_registerName.restype = objc_id
-libmetal.MTLCreateSystemDefaultDevice.restype = objc_instance
-libdispatch.dispatch_data_create.restype = objc_instance
 
+# FIXME: these need autogen to support objc categories
+# https://developer.apple.com/library/archive/documentation/Cocoa/Conceptual/ObjectiveC/Chapters/ocCategories.html
 @functools.cache
-def msg(selector: str, restype: type[T] = objc_id):  # type: ignore [assignment]
-  resname = libobjc.sel_registerName(selector.encode())
-  sender = libobjc["objc_msgSend"] # Using attribute access returns a new reference so setting restype is safe
-  sender.restype = restype
-  def _msg(ptr: objc_id, *args: Any) -> T: return sender(ptr, resname, *args)
-  return _msg
-
-@functools.cache
-def to_ns_str(s: str): return msg("stringWithUTF8String:", objc_instance)(libobjc.objc_getClass(b"NSString"), s.encode())
-def from_ns_str(s): return bytes(msg("UTF8String", ctypes.c_char_p)(s)).decode()
+def to_ns_str(s: str): return ctypes.cast(objc.msg("stringWithUTF8String:")(metal.NSString._objc_class_, s.encode()), metal.NSString)
+def from_ns_str(s): return bytes(objc.msg("UTF8String", ctypes.c_char_p)(s)).decode()
 
 def to_struct(*t: int, _type: type[ctypes._SimpleCData] = ctypes.c_ulong):
   return init_c_struct_t(tuple([(f"field{i}", _type) for i in range(len(t))]))(*t)
 
-def wait_check(cbuf: Any):
-  msg("waitUntilCompleted")(cbuf)
-  error_check(msg("error", objc_instance)(cbuf))
+def wait_check(cbuf:metal.MTLCommandBuffer):
+  cbuf.waitUntilCompleted()
+  error_check(cbuf.error().retained())
 
-def cmdbuf_label(cbuf: objc_id) -> str|None: return from_ns_str(label) if (label:=msg("label", objc_id)(cbuf)).value is not None else None
-def cmdbuf_st_time(cbuf: objc_id) -> float: return cast(float, msg("GPUStartTime", ctypes.c_double)(cbuf))
-def cmdbuf_en_time(cbuf: objc_id) -> float: return cast(float, msg("GPUEndTime", ctypes.c_double)(cbuf))
+def cmdbuf_label(cbuf:metal.MTLCommandBuffer) -> str|None: return from_ns_str(label) if (label:=cbuf.label()).value is not None else None
 
-def error_check(error: objc_instance, error_constructor: type[Exception] = RuntimeError):
+def error_check(error: metal.NSError, error_constructor: type[Exception] = RuntimeError):
   if error.value is None: return None
-  raise error_constructor(from_ns_str(msg("localizedDescription", objc_instance)(error)))
+  raise error_constructor(from_ns_str(error.localizedDescription().retained()))
 
 class MetalDevice(Compiled):
   def __init__(self, device:str):
-    self.sysdevice = libmetal.MTLCreateSystemDefaultDevice()
-    self.mtl_queue = msg("newCommandQueueWithMaxCommandBufferCount:", objc_instance)(self.sysdevice, 1024)
+    self.sysdevice = metal.MTLCreateSystemDefaultDevice()
+    self.mtl_queue = self.sysdevice.newCommandQueueWithMaxCommandBufferCount(1024)
     if self.mtl_queue is None: raise RuntimeError("Cannot allocate a new command queue")
-    self.mtl_buffers_in_flight: list[Any] = []
-    self.timeline_signal = msg("newSharedEvent", objc_instance)(self.sysdevice)
+    self.mtl_buffers_in_flight: list[metal.MTLCommandBuffer] = []
+    self.timeline_signal = self.sysdevice.newSharedEvent()
     self.timeline_value = 0
 
     Compiled.profile_events += [ProfileDeviceEvent(device)]
@@ -76,23 +44,22 @@ class MetalDevice(Compiled):
     from tinygrad.runtime.graph.metal import MetalGraph
     # NOTE: GitHub CI macOS runners use paravirtualized metal which is broken with graph.
     # This can be reproduced locally with any virtualization software (like utm) that can create macOS VMs with apple's own virtualization framework.
-    super().__init__(device, MetalAllocator(self), [(MetalRenderer, MetalCompiler), (MetalRenderer, Compiler)],
-                     functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(msg('name')(self.sysdevice)).lower() else None)
+    super().__init__(device, MetalAllocator(self), [(MetalRenderer, MetalCompiler), (MetalRenderer, Compiler)], functools.partial(MetalProgram, self),
+                     MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None)
 
   def synchronize(self):
     for cbuf in self.mtl_buffers_in_flight:
       wait_check(cbuf)
-      st, en = decimal.Decimal(cmdbuf_st_time(cbuf)) * 1000000, decimal.Decimal(cmdbuf_en_time(cbuf)) * 1000000
+      st, en = decimal.Decimal(cbuf.GPUStartTime()) * 1000000, decimal.Decimal(cbuf.GPUEndTime()) * 1000000
       # NOTE: command buffers from MetalGraph are not profiled here
       if PROFILE and (lb:=cmdbuf_label(cbuf)) is not None and not lb.startswith("batched"):
         Compiled.profile_events += [ProfileRangeEvent(self.device, lb, st, en, is_copy=lb.startswith("COPY"))]
     self.mtl_buffers_in_flight.clear()
 
-def metal_src_to_library(device:MetalDevice, src:str) -> objc_instance:
-  options = msg("new", objc_instance)(libobjc.objc_getClass(b"MTLCompileOptions"))
-  msg("setFastMathEnabled:")(options, getenv("METAL_FAST_MATH"))
-  library = msg("newLibraryWithSource:options:error:", objc_instance)(device.sysdevice, to_ns_str(src),
-                                                                      options, ctypes.byref(compileError:=objc_instance()))
+def metal_src_to_library(device:MetalDevice, src:str) -> metal.MTLLibrary:
+  options = metal.MTLCompileOptions.new()
+  options.setFastMathEnabled(getenv("METAL_FAST_MATH"))
+  library = device.sysdevice.newLibraryWithSource_options_error(to_ns_str(src), options, ctypes.byref(compileError:=metal.NSError().retained()))
   error_check(compileError, CompileError)
   return library
 
@@ -155,71 +122,71 @@ class MetalProgram:
     self.dev, self.name, self.lib = dev, name, lib
     if lib[:4] == b"MTLB":
       # binary metal library
-      data = libdispatch.dispatch_data_create(lib, len(lib), None, None)
-      self.library = msg("newLibraryWithData:error:", objc_instance)(self.dev.sysdevice, data, ctypes.byref(error_lib:=objc_instance()))
+      data = objc.dispatch_data_create(lib, len(lib), None, None)
+      self.library = self.dev.sysdevice.newLibraryWithData_error(data, ctypes.byref(error_lib:=metal.NSError().retained())).retained()
       error_check(error_lib)
     else:
       # metal source. rely on OS caching
       try: self.library = metal_src_to_library(self.dev, lib.decode())
       except CompileError as e: raise RuntimeError from e
-    self.fxn = msg("newFunctionWithName:", objc_instance)(self.library, to_ns_str(name))
-    descriptor = msg("new", objc_instance)(libobjc.objc_getClass(b"MTLComputePipelineDescriptor"))
-    msg("setComputeFunction:")(descriptor, self.fxn)
-    msg("setSupportIndirectCommandBuffers:")(descriptor, True)
-    self.pipeline_state = msg("newComputePipelineStateWithDescriptor:options:reflection:error:", objc_instance)(self.dev.sysdevice,
-      descriptor, MTLPipelineOption.MTLPipelineOptionNone, None, ctypes.byref(error_pipeline_creation:=objc_instance()))
+    self.fxn = self.library.newFunctionWithName(to_ns_str(name)).retained()
+    descriptor = metal.MTLComputePipelineDescriptor.new()
+    descriptor.setComputeFunction(self.fxn)
+    descriptor.setSupportIndirectCommandBuffers(True)
+    self.pipeline_state = self.dev.sysdevice.newComputePipelineStateWithDescriptor_options_reflection_error(descriptor, metal.MTLPipelineOptionNone,
+      None, ctypes.byref(error_pipeline_creation:=metal.NSError().retained()))
     error_check(error_pipeline_creation)
     # cache these msg calls
-    self.max_total_threads: int = cast(int, msg("maxTotalThreadsPerThreadgroup", ctypes.c_ulong)(self.pipeline_state))
+    self.max_total_threads: int = self.pipeline_state.maxTotalThreadsPerThreadgroup()
 
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
     if prod(local_size) > self.max_total_threads:
-      exec_width = msg("threadExecutionWidth", ctypes.c_ulong)(self.pipeline_state)
-      memory_length = msg("staticThreadgroupMemoryLength", ctypes.c_ulong)(self.pipeline_state)
+      exec_width = self.pipeline_state.threadExecutionWidth()
+      memory_length = self.pipeline_state.staticThreadgroupMemoryLength()
       raise RuntimeError(f"local size {local_size} bigger than {self.max_total_threads} with exec width {exec_width} memory length {memory_length}")
-    command_buffer = msg("commandBuffer", objc_instance)(self.dev.mtl_queue)
-    encoder = msg("computeCommandEncoder", objc_instance)(command_buffer)
-    msg("setComputePipelineState:")(encoder, self.pipeline_state)
-    for i,a in enumerate(bufs): msg("setBuffer:offset:atIndex:")(encoder, a.buf, a.offset, i)
-    for i,a in enumerate(vals, start=len(bufs)): msg("setBytes:length:atIndex:")(encoder, bytes(ctypes.c_int(a)), 4, i)
-    msg("dispatchThreadgroups:threadsPerThreadgroup:")(encoder, to_struct(*global_size), to_struct(*local_size))
-    msg("endEncoding")(encoder)
-    msg("setLabel:")(command_buffer, to_ns_str(self.name)) # TODO: is this always needed?
-    msg("commit")(command_buffer)
+    command_buffer = self.dev.mtl_queue.commandBuffer().retained() # FIXME: is this really ARC?
+    encoder = command_buffer.computeCommandEncoder().retained() # FIXME: is this really ARC?
+    encoder.setComputePipelineState(self.pipeline_state)
+    for i,a in enumerate(bufs): encoder.setBuffer_offset_atIndex(a.buf, a.offset, i)
+    for i,a in enumerate(vals, start=len(bufs)): encoder.setBytes_length_atIndex(bytes(ctypes.c_int(a)), 4, i)
+    encoder.dispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_size), metal.MTLSize(*local_size))
+    encoder.endEncoding()
+    command_buffer.setLabel(to_ns_str(self.name)) # TODO: is this always needed?
+    command_buffer.commit()
     self.dev.mtl_buffers_in_flight.append(command_buffer)
     if wait:
       wait_check(command_buffer)
-      return cmdbuf_en_time(command_buffer) - cmdbuf_st_time(command_buffer)
+      return command_buffer.GPUEndTime() - command_buffer.GPUStartTime()
 
 class MetalBuffer:
-  def __init__(self, buf:Any, size:int, offset=0): self.buf, self.size, self.offset = buf, size, offset
+  def __init__(self, buf:metal.MTLBuffer, size:int, offset=0): self.buf, self.size, self.offset = buf, size, offset
 
 class MetalAllocator(LRUAllocator[MetalDevice]):
   def _alloc(self, size:int, options) -> MetalBuffer:
-    if options.external_ptr: return MetalBuffer(objc_id(options.external_ptr), size)
+    if options.external_ptr: return MetalBuffer(metal.MTLBuffer(options.external_ptr), size)
 
     # Buffer is explicitly released in _free() rather than garbage collected via reference count
-    ret = msg("newBufferWithLength:options:", objc_id)(self.dev.sysdevice, ctypes.c_ulong(size), MTLResourceOptions.MTLResourceStorageModeShared)
+    ret = self.dev.sysdevice.newBufferWithLength_options(size, metal.MTLResourceStorageModeShared)
+    ret.retain = False
     if ret.value is None: raise MemoryError(f"Metal OOM while allocating {size=}")
     return MetalBuffer(ret, size)
   def _free(self, opaque:MetalBuffer, options):
-    if msg is not None and libobjc is not None: msg("release")(opaque.buf)
+    if not sys.is_finalizing(): opaque.buf.release
   def _transfer(self, dest:MetalBuffer, src:MetalBuffer, sz:int, src_dev:MetalDevice, dest_dev:MetalDevice):
     dest_dev.synchronize()
-    src_command_buffer = msg("commandBuffer", objc_instance)(src_dev.mtl_queue)
-    encoder = msg("blitCommandEncoder", objc_instance)(src_command_buffer)
-    msg("copyFromBuffer:sourceOffset:toBuffer:destinationOffset:size:")(encoder, src.buf, ctypes.c_ulong(src.offset),
-        dest.buf, ctypes.c_ulong(dest.offset), ctypes.c_ulong(sz))
-    msg("endEncoding")(encoder)
+    src_command_buffer = src_dev.mtl_queue.commandBuffer().retained()
+    encoder = src_command_buffer.blitCommandEncoder().retained()
+    encoder.copyFromBuffer_sourceOffset_toBuffer_destinationOffset_size(src.buf, src.offset, dest.buf, dest.offset, sz)
+    encoder.endEncoding()
     if src_dev != dest_dev:
-      msg("encodeSignalEvent:value:")(src_command_buffer, src_dev.timeline_signal, src_dev.timeline_value)
-      dest_command_buffer = msg("commandBuffer", objc_instance)(dest_dev.mtl_queue)
-      msg("encodeWaitForEvent:value:")(dest_command_buffer, src_dev.timeline_signal, src_dev.timeline_value)
-      msg("commit")(dest_command_buffer)
+      src_command_buffer.encodeSignalEvent_value(ctypes.cast(src_dev.timeline_signal, metal.MTLEvent), src_dev.timeline_value)
+      dest_command_buffer = dest_dev.mtl_queue.commandBuffer().retained()
+      dest_command_buffer.encodeWaitForEvent_value(ctypes.cast(src_dev.timeline_signal, metal.MTLEvent), src_dev.timeline_value)
+      dest_command_buffer.commit()
       dest_dev.mtl_buffers_in_flight.append(dest_command_buffer)
       src_dev.timeline_value += 1
-    msg("setLabel:")(src_command_buffer, to_ns_str(f"COPY {src_dev.device} -> {dest_dev.device}"))
-    msg("commit")(src_command_buffer)
+    src_command_buffer.setLabel(to_ns_str(f"COPY {src_dev.device} -> {dest_dev.device}"))
+    src_command_buffer.commit()
     src_dev.mtl_buffers_in_flight.append(src_command_buffer)
     # Transfers currently synchronize the completion. Otherwise, copies can sometimes lead to incorrect values.
     # There is no real metal multidevice support for now, so transfer is used only for tests.
@@ -228,7 +195,7 @@ class MetalAllocator(LRUAllocator[MetalDevice]):
     with cpu_profile(prof_desc, self.dev.device, is_copy=True): dst[:] = src
   def _as_buffer(self, src:MetalBuffer) -> memoryview:
     self.dev.synchronize()
-    return to_mv(cast(int, msg("contents", objc_id)(src.buf).value), src.size + src.offset)[src.offset:]
+    return to_mv(src.buf.contents(), src.size + src.offset)[src.offset:]
   def _copyin(self, dest:MetalBuffer, src:memoryview): self._cp_mv(self._as_buffer(dest), src, "TINY -> METAL")
   def _copyout(self, dest:memoryview, src:MetalBuffer): self._cp_mv(dest, self._as_buffer(src), "METAL -> TINY")
   def _offset(self, buf:MetalBuffer, size:int, offset:int): return MetalBuffer(buf.buf, size, offset)

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,115 +1,268 @@
-import ctypes.util, importlib.metadata, itertools, re, functools, os
+import ctypes, itertools, re, functools, os
 from tinygrad.helpers import flatten, unwrap
-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
+from tinygrad.runtime.autogen import libclang as clang # hmmm
+from typing import Iterator
 
-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 unwrap_cursor(c: clang.CXCursor) -> clang.CXCursor:
+  assert c != clang.clang_getNullCursor()
+  return c
+
+def children(c: clang.CXCursor) -> list[clang.CXCursor]:
+  ret = []
+  @clang.CXCursorVisitor
+  def visitor(child, _0, _1):
+    nonlocal ret
+    ret.append(child)
+    return clang.CXChildVisit_Continue
+  clang.clang_visitChildren(c, visitor, None)
+  return ret
+
+def fields(t: clang.CXType) -> list[clang.CXCursor]:
+  ret = []
+  @clang.CXFieldVisitor
+  def visitor(child, _):
+    nonlocal ret
+    ret.append(child)
+    return clang.CXVisit_Continue
+  clang.clang_Type_visitFields(t, visitor, None)
+  return ret
+
+def walk(c: clang.CXCursor) -> Iterator[clang.CXCursor]:
+  yield c
+  for child in children(c): yield from walk(child)
+
+def arguments(c: clang.CXCursor|clang.CXType):
+  yield from ((clang.clang_Cursor_getArgument if isinstance(c, clang.CXCursor) else clang.clang_getArgType)(c, i)
+              for i in range(clang.clang_Cursor_getNumArguments(c) if isinstance(c, clang.CXCursor) else clang.clang_getNumArgTypes(c)))
+
+class Tokens:
+  def __init__(self, c: clang.CXCursor):
+    clang.clang_tokenize(tu:=clang.clang_Cursor_getTranslationUnit(c), clang.clang_getCursorExtent(c),
+                         toks:=(ctypes.POINTER(clang.CXToken)()), cnt:=ctypes.c_uint32())
+    self.tu, self.toks = tu, toks[:cnt.value]
+    for t in self.toks: t._tu = tu
+
+  def __getitem__(self, idx): return self.toks[idx]
+  def __len__(self): return len(self.toks)
+
+  def __del__(self):
+    if self.toks: clang.clang_disposeTokens(self.tu, self.toks[0], len(self.toks))
+
+def cxs(fn):
+  @functools.wraps(fn)
+  def wrap(*args, **kwargs) -> str:
+    if ctypes.cast(clang.clang_getCString(cxs:=fn(*args, **kwargs)), ctypes.c_void_p).value is None: return ""
+    ret = ctypes.string_at(clang.clang_getCString(cxs)).decode()
+    clang.clang_disposeString(cxs)
+    return ret
+  return wrap
+
+# TODO: caching this would be nice?
+@cxs
+def nm(c: clang.CXCursor|clang.CXToken|clang.CXType) -> str:
+  return clang.clang_getTokenSpelling(c._tu, c) if isinstance(c, clang.CXToken) else getattr(clang, f"clang_get{c.__class__.__name__[2:]}Spelling")(c)
+def extent(c: clang.CXCursor|clang.CXToken|clang.CXType) -> clang.CXSourceRange:
+  return clang.clang_getTokenExtent(c._tu, c) if isinstance(c, clang.CXToken) else getattr(clang, f"clang_get{c.__class__.__name__[2:]}Extent")(c)
+def loc(c: clang.CXCursor|clang.CXToken|clang.CXType) -> clang.CXSourceLocation:
+  return clang.clang_getTokenLocation(c._tu, c) if isinstance(c, clang.CXToken) else getattr(clang, f"clang_get{c.__class__.__name__[2:]}Location")(c)
+@cxs
+def loc_file(loc: clang.CXSourceLocation) -> str:
+  clang.clang_getExpansionLocation(loc, f:=clang.CXFile(), None, None, None)
+  return clang.clang_getFileName(f)
+def loc_off(loc: clang.CXSourceLocation) -> int:
+  clang.clang_getExpansionLocation(loc, None, None, None, off:=ctypes.c_uint32())
+  return off.value
+def loc_line(loc: clang.CXSourceLocation) -> int:
+  clang.clang_getExpansionLocation(loc, None, line:=ctypes.c_uint32(), None, None)
+  return line.value
 
-def fst(c): return next(c.get_children())
-def last(c): return list(c.get_children())[-1]
 def readext(f, fst, snd=None):
-  with open(f, "r") as f:
-    f.seek(start:=(fst.start.offset if isinstance(fst, SourceRange) else fst))
-    return f.read((fst.end.offset if isinstance(fst, SourceRange) else snd)-start)
-def attrs(c): return list(filter(lambda k: (v:=k.value) >= 400 and v < 500, map(lambda c: c.kind, c.get_children())))
+  with open(f, "r") as f: # reopening this every time is dumb...
+    f.seek(start:=loc_off(clang.clang_getRangeStart(fst) if isinstance(fst, clang.CXSourceRange) else fst))
+    return f.read(loc_off(clang.clang_getRangeEnd(fst) if isinstance(fst, clang.CXSourceRange) else snd)-start)
+def attrs(c): return list(filter(lambda k: (v:=k.value) >= 400 and v < 500, map(lambda c: c.kind, children(c))))
+
+def protocols(t): yield from (clang.clang_Type_getObjCProtocolDecl(t, i) for i in range(clang.clang_Type_getNumObjCProtocolRefs(t)))
+def basetype(t): return clang.clang_Type_getObjCObjectBaseType(t)
 
 base_rules = [(r'\s*\\\n\s*', ' '), (r'\s*\n\s*', ' '), (r'//.*', ''), (r'/\*.*?\*/', ''), (r'\b(0[xX][0-9a-fA-F]+|\d+)[uUlL]+\b', r'\1'),
               (r'\b0+(?=\d)', ''), (r'\s*&&\s*', r' and '), (r'\s*\|\|\s*', r' or '), (r'\s*!\s*', ' not '),
               (r'(struct|union|enum)\s*([a-zA-Z_][a-zA-Z0-9_]*\b)', r'\1_\2'),
               (r'\((unsigned )?(char|uint64_t)\)', ''), (r'^.*\d+:\d+.*$', ''), (r'^.*\w##\w.*$', '')]
 
-ints = (TK.INT, TK.UINT, TK.LONG, TK.ULONG, TK.LONGLONG, TK.ULONGLONG)
+uints = (clang.CXType_Char_U, clang.CXType_UChar, clang.CXType_UShort, clang.CXType_UInt, clang.CXType_ULong, clang.CXType_ULongLong)
+ints = uints + (clang.CXType_Char_S, clang.CXType_Short, clang.CXType_Int, clang.CXType_ULong, clang.CXType_LongLong)
+fns = (clang.CXType_FunctionProto, clang.CXType_FunctionNoProto)
+specs = (clang.CXCursor_ObjCSuperClassRef,)
+# https://clang.llvm.org/docs/AutomaticReferenceCounting.html#arc-method-families
+arc_families = ['alloc', 'copy', 'mutableCopy', 'new']
 
 def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_errno=False, anon_names={}, types={}, parse_macros=True):
-  macros, lines, anoncnt, types = [], [], itertools.count().__next__, {k:(v,True) for k,v in types.items()}
+  macros, lines, anoncnt, types, objc = [], [], itertools.count().__next__, {k:(v,True) for k,v in types.items()}, False
   def tname(t, suggested_name=None, typedef=None) -> str:
-    suggested_name = anon_names.get(f"{(decl:=t.get_declaration()).location.file}:{decl.location.line}", suggested_name)
-    nonlocal lines, types, anoncnt
-    tmap = {TK.VOID:"None", TK.CHAR_U:"ctypes.c_ubyte", TK.UCHAR:"ctypes.c_ubyte", TK.CHAR_S:"ctypes.c_char", TK.SCHAR:"ctypes.c_char",
-            **{getattr(TK, k):f"ctypes.c_{k.lower()}" for k in ["BOOL", "WCHAR", "FLOAT", "DOUBLE", "LONGDOUBLE"]},
-            **{getattr(TK, k):f"ctypes.c_{'u' if 'U' in k else ''}int{sz}" for sz,k in
-               [(16, "USHORT"), (16, "SHORT"), (32, "UINT"), (32, "INT"), (64, "ULONG"), (64, "LONG"), (64, "ULONGLONG"), (64, "LONGLONG")]}}
+    suggested_name = anon_names.get(f"{loc_file(loc(decl:=clang.clang_getTypeDeclaration(t)))}:{loc_line(loc(decl))}", suggested_name)
+    nonlocal lines, types, anoncnt, objc
+    tmap = {clang.CXType_Void:"None", clang.CXType_Char_U:"ctypes.c_ubyte", clang.CXType_UChar:"ctypes.c_ubyte", clang.CXType_Char_S:"ctypes.c_char",
+            clang.CXType_SChar:"ctypes.c_char",
+            **{getattr(clang, f'CXType_{k}'):f"ctypes.c_{k.lower()}" for k in ["Bool", "WChar", "Float", "Double", "LongDouble"]},
+            **{getattr(clang, f'CXType_{k}'):f"ctypes.c_{'u' if 'U' in k else ''}int{sz}" for sz,k in
+               [(16, "UShort"), (16, "Short"), (32, "UInt"), (32, "Int"), (64, "ULong"), (64, "Long"), (64, "ULongLong"), (64, "LongLong")]}}
 
     if t.kind in tmap: return tmap[t.kind]
-    if t.spelling in types and types[t.spelling][1]: return types[t.spelling][0]
-    if ((f:=t).kind in (fks:=(TK.FUNCTIONPROTO, TK.FUNCTIONNOPROTO))) or (t.kind == TK.POINTER and (f:=t.get_pointee()).kind in fks):
-      return f"ctypes.CFUNCTYPE({tname(f.get_result())}{(', '+', '.join(map(tname, f.argument_types()))) if f.kind==TK.FUNCTIONPROTO else ''})"
+    if nm(t) in types and types[nm(t)][1]: return types[nm(t)][0]
+    if ((f:=t).kind in fns) or (t.kind == clang.CXType_Pointer and (f:=clang.clang_getPointeeType(t)).kind in fns):
+      return (f"ctypes.CFUNCTYPE({tname(clang.clang_getResultType(f))}" +
+              ((', '+', '.join(map(tname, arguments(f)))) if f.kind==clang.CXType_FunctionProto else '') + ")")
     match t.kind:
-      case TK.POINTER: return "ctypes.c_void_p" if (ptr:=t.get_pointee()).kind == TK.VOID else f"ctypes.POINTER({tname(ptr)})"
-      case TK.ELABORATED: return tname(t.get_named_type(), suggested_name)
-      case TK.TYPEDEF if t.spelling == t.get_canonical().spelling: return tname(t.get_canonical())
-      case TK.TYPEDEF:
-        defined, nm = (canon:=t.get_canonical()).spelling in types, tname(canon, typedef=t.spelling.replace('::', '_'))
-        types[t.spelling] = nm if t.spelling.startswith("__") else t.spelling.replace('::', '_'), True
+      case clang.CXType_Pointer:
+        return "ctypes.c_void_p" if (p:=clang.clang_getPointeeType(t)).kind==clang.CXType_Void else f"ctypes.POINTER({tname(p)})"
+      case clang.CXType_ObjCObjectPointer: return tname(clang.clang_getPointeeType(t)) # TODO: this seems wrong
+      case clang.CXType_Elaborated: return tname(clang.clang_Type_getNamedType(t), suggested_name)
+      case clang.CXType_Typedef if nm(t) == nm(canon:=clang.clang_getCanonicalType(t)): return tname(canon)
+      case clang.CXType_Typedef:
+        defined, cnm = nm(canon:=clang.clang_getCanonicalType(t)) in types, tname(canon, typedef=nm(t).replace('::', '_'))
+        types[nm(t)] = cnm if nm(t).startswith("__") else nm(t).replace('::', '_'), True
         # RECORDs need to handle typedefs specially to allow for self-reference
-        if canon.kind != TK.RECORD or defined: lines.append(f"{t.spelling.replace('::', '_')} = {nm}")
-        return types[t.spelling][0]
-      case TK.RECORD:
+        if canon.kind != clang.CXType_Record or defined: lines.append(f"{nm(t).replace('::', '_')} = {cnm}")
+        return types[nm(t)][0]
+      case clang.CXType_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
+        if nm(t) in types: types[nm(t)] = (tnm:=types[nm(t)][0]), len(fields(t)) != 0
         else:
-          if decl.is_anonymous():
-            types[t.spelling] = (nm:=(suggested_name or (f"_anon{'struct' if decl.kind == CK.STRUCT_DECL else 'union'}{anoncnt()}")), True)
-          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 clang.clang_Cursor_isAnonymous(decl):
+            types[nm(t)] = (tnm:=(suggested_name or (f"_anon{'struct' if decl.kind==clang.CXCursor_StructDecl else 'union'}{anoncnt()}")), True)
+          else: types[nm(t)] = (tnm:=nm(t).replace(' ', '_').replace('::', '_')), len(fields(t)) != 0
+          lines.append(f"class {tnm}({'Struct' if decl.kind==clang.CXCursor_StructDecl else 'ctypes.Union'}): pass")
+          if typedef: lines.append(f"{typedef} = {tnm}")
+        if ((is_packed:=(clang.CXCursor_PackedAttr in attrs(decl)) or
+            ((N:=clang.clang_Type_getAlignOf(t)) != max([clang.clang_Type_getAlignOf(clang.clang_getCursorType(f)) for f in fields(t)], default=N)))):
+          if clang.clang_Type_getAlignOf(t) != 1:
+            print(f"WARNING: ignoring alignment={clang.clang_Type_getAlignOf(t)} on {nm(t)}")
+            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 []))
-        return nm
-      case TK.ENUM:
+        def is_anon(f): return clang.clang_Cursor_isAnonymousRecordDecl(clang.clang_getTypeDeclaration(clang.clang_getCursorType(f)))
+        ll=["  ("+((fn:=f"'_{acnt()}'")+f", {tname(clang.clang_getCursorType(f), tnm+fn[1:-1])}" if is_anon(f) else f"'{nm(f)}', "+
+            tname(clang.clang_getCursorType(f), f'{tnm}_{nm(f)}'))+(f',{clang.clang_getFieldDeclBitWidth(f)}' * clang.clang_Cursor_isBitField(f))+"),"
+            for f in fields(t)]
+        lines.extend(([f"{tnm}._anonymous_ = ["+", ".join(f"'_{i}'" for i in range(n))+"]"] if (n:=acnt()) else [])+
+                     ([f"{tnm}._packed_ = True"] * is_packed)+([f"{tnm}._fields_ = [",*ll,"]"] if ll else []))
+        return tnm
+      case clang.CXType_Enum:
         # TODO: C++ and GNU C have forward declared enums
-        if decl.is_anonymous(): types[t.spelling] = suggested_name or f"_anonenum{anoncnt()}", True
-        else: types[t.spelling] = t.spelling.replace(' ', '_').replace('::', '_'), True
-        lines.append(f"{types[t.spelling][0]} = CEnum({tname(decl.enum_type)})\n" +
-                     "\n".join(f"{e.spelling} = {types[t.spelling][0]}.define('{e.spelling}', {e.enum_value})" for e in decl.get_children()
-                     if e.kind == CK.ENUM_CONSTANT_DECL) + "\n")
-        return types[t.spelling][0]
-      case TK.CONSTANTARRAY:
-        return f"({tname(t.get_array_element_type(), suggested_name.rstrip('s') if suggested_name else None)} * {t.get_array_size()})"
-      case TK.INCOMPLETEARRAY: return f"({tname(t.get_array_element_type(), suggested_name.rstrip('s') if suggested_name else None)} * 0)"
+        if clang.clang_Cursor_isAnonymous(decl): types[nm(t)] = suggested_name or f"_anonenum{anoncnt()}", True
+        else: types[nm(t)] = nm(t).replace(' ', '_').replace('::', '_'), True
+        ety = clang.clang_getEnumDeclIntegerType(decl)
+        def value(e): return (clang.clang_getEnumConstantDeclUnsignedValue if ety.kind in uints else clang.clang_getEnumConstantDeclValue)(e)
+        lines.append(f"{types[nm(t)][0]} = CEnum({tname(ety)})\n" +
+                     "\n".join(f"{nm(e)} = {types[nm(t)][0]}.define('{nm(e)}', {value(e)})" for e in children(decl)
+                     if e.kind == clang.CXCursor_EnumConstantDecl) + "\n")
+        return types[nm(t)][0]
+      case clang.CXType_ConstantArray:
+        return f"({tname(clang.clang_getArrayElementType(t),suggested_name.rstrip('s') if suggested_name else None)} * {clang.clang_getArraySize(t)})"
+      case clang.CXType_IncompleteArray:
+        return f"({tname(clang.clang_getArrayElementType(t), suggested_name.rstrip('s') if suggested_name else None)} * 0)"
+      case clang.CXType_ObjCInterface:
+        is_defn = bool([f.kind for f in children(decl) if f.kind in (clang.CXCursor_ObjCInstanceMethodDecl, clang.CXCursor_ObjCClassMethodDecl)])
+        if (tnm:=nm(t)) not in types: lines.append(f"class {tnm}(objc.Spec): pass")
+        types[tnm] = tnm, is_defn
+        if is_defn:
+          ims, cms = parse_objc_spec(decl, tnm, clang.CXCursor_ObjCInstanceMethodDecl), parse_objc_spec(decl, tnm, clang.CXCursor_ObjCClassMethodDecl)
+          bases = [tname(clang.clang_getCursorType(b)) for b in children(decl) if b.kind in specs]
+          lines.extend([*([f"{tnm}._bases_ = [{', '.join(bases)}]"] if bases else []),
+                        *([f"{tnm}._methods_ = [", *ims, ']'] if ims else []), *([f"{tnm}._classmethods_ = [", *cms, ']'] if cms else [])])
+        return tnm
+      case clang.CXType_ObjCSel: return "objc.id_"
+      case clang.CXType_ObjCId: return (objc:=True, "objc.id_")[1]
+      case clang.CXType_ObjCObject:
+        if basetype(t).kind != clang.CXType_ObjCId: raise NotImplementedError(f"generics unsupported: {nm(t)}")
+        if len(ps:=[proto(p) for p in protocols(t)]) == 0:
+          types[nm(t)] = "objc.id_", True
+          return "objc.id_"
+        if len(ps) == 1:
+          types[nm(t)] = ps[0], True
+          return ps[0]
+        types[nm(t)] = (tnm:=f"_anondynamic{anoncnt()}"), True
+        lines.append(f"class {tnm}({', '.join(ps)}): pass # {nm(t)}")
+        return tnm
       case _: raise NotImplementedError(f"unsupported type {t.kind}")
 
+  # parses an objc @interface or @protocol, returning a list of declerations that objc.Spec can parse, for the specified kind
+  # NB: ivars are unsupported
+  def parse_objc_spec(decl:clang.CXCursor, dnm:str, kind) -> list[str]:
+    nonlocal lines, types
+    ms = []
+    for d in filter(lambda d: d.kind == kind, children(decl)):
+      rollback = lines, types
+      try: ms.append(f"  ('{nm(d)}', {repr('instancetype') if nm(rt:=clang.clang_getCursorResultType(d))=='instancetype' else tname(rt)}, " +
+        f"[{', '.join('instancetype' if nm(a) == 'instancetype' else tname(clang.clang_getCursorType(a)) for a in arguments(d))}]" +
+        (", True" * (clang.CXCursor_NSReturnsRetained in attrs(d) or (any(nm(d).startswith(s) for s in arc_families) and rt.kind!=clang.CXType_Void)))
+                     + "),")
+      except NotImplementedError as e:
+        print(f"skipping {dnm}.{nm(d)}: {e}")
+        lines, types = rollback
+    return ms
+
+  # libclang doesn't have a "type" for @protocol, so we have to do this here...
+  def proto(decl):
+    nonlocal lines, types
+    if (dnm:=nm(decl)) in types and types[dnm][1]: return types[dnm][0]
+    # check if this is a forward declaration
+    is_defn = bool([f.kind for f in children(decl) if f.kind in (clang.CXCursor_ObjCInstanceMethodDecl, clang.CXCursor_ObjCClassMethodDecl)])
+    if dnm not in types: lines.append(f"class {dnm}(objc.Spec): pass")
+    types[dnm] = dnm, is_defn
+    if is_defn:
+      bases = [proto(b) for b in children(decl) if b.kind==clang.CXCursor_ObjCProtocolRef and nm(b) != nm(decl)]
+      ims, cms = parse_objc_spec(decl, dnm, clang.CXCursor_ObjCInstanceMethodDecl), parse_objc_spec(decl, dnm, clang.CXCursor_ObjCClassMethodDecl)
+      lines.extend([*([f"{dnm}._bases_ = [{', '.join(bases)}]"] if bases else []),
+                    *([f"{dnm}._methods_ = [", *ims, "]"] if ims else []), *([f"{dnm}._classmethods_ = [", *cms, "]"] if cms else [])])
+    return dnm
+
   for f in files:
-    tu = Index.create().parse(f, args, options=TU.PARSE_DETAILED_PROCESSING_RECORD)
-    (pp:=PP.create(tu.cursor)).set_property(PPP.TerseOutput, 1)
-    for c in tu.cursor.walk_preorder():
-      if str(c.location.file) != str(f) and (not recsym or c.kind not in (CK.FUNCTION_DECL,)): continue
+    aa = ctypes.cast((ctypes.c_char_p * len(args))(*[x.encode() for x in args]), ctypes.POINTER(ctypes.POINTER(ctypes.c_char))) if len(args) else None
+    tu = clang.clang_parseTranslationUnit(idx:=clang.clang_createIndex(False, 0), os.fspath(f).encode(), aa, len(args), None, 0,
+                                          clang.CXTranslationUnit_DetailedPreprocessingRecord)
+    # FIXME: deep walk is not neccesary...
+    for c in walk(unwrap_cursor(clang.clang_getTranslationUnitCursor(tu))):
+      if loc_file(loc(c)) != str(f) and (not recsym or c.kind not in (clang.CXCursor_FunctionDecl,)):
+        continue
       rollback = lines, types
       try:
         match c.kind:
-          case CK.FUNCTION_DECL if c.linkage == LK.EXTERNAL and dll:
+          case clang.CXCursor_FunctionDecl if clang.clang_getCursorLinkage(c) == clang.CXLinkage_External and dll:
             # TODO: we could support name-mangling
-            lines.append(f"# {c.pretty_printed(pp)}\ntry: ({c.spelling}:=dll.{c.spelling}).restype, {c.spelling}.argtypes = "
-              f"{tname(c.result_type)}, [{', '.join(tname(arg.type) for arg in c.get_arguments())}]\nexcept AttributeError: pass\n")
-          case CK.STRUCT_DECL | CK.UNION_DECL | CK.TYPEDEF_DECL | CK.ENUM_DECL: tname(c.type)
-          case CK.MACRO_DEFINITION if parse_macros and len(toks:=list(c.get_tokens())) > 1:
-            if toks[1].spelling == '(' and toks[0].extent.end.column == toks[1].extent.start.column:
+            lines.append(f"try: ({nm(c)}:=dll.{nm(c)}).restype, {nm(c)}.argtypes = {tname(clang.clang_getCursorResultType(c))}, "
+                         f"[{', '.join(tname(clang.clang_getCursorType(arg)) for arg in arguments(c))}]\nexcept AttributeError: pass\n")
+            if clang.CXCursor_NSReturnsRetained in attrs(c): lines.append(f"{nm(c)} = objc.returns_retained({nm(c)})")
+          case (clang.CXCursor_StructDecl | clang.CXCursor_UnionDecl | clang.CXCursor_TypedefDecl | clang.CXCursor_EnumDecl
+                | clang.CXCursor_ObjCInterfaceDecl): tname(clang.clang_getCursorType(c))
+          case clang.CXCursor_MacroDefinition if parse_macros and len(toks:=Tokens(c)) > 1:
+            if nm(toks[1])=='(' and clang.clang_equalLocations(clang.clang_getRangeEnd(extent(toks[0])), clang.clang_getRangeStart(extent(toks[1]))):
               it = iter(toks[1:])
-              _args = [t.spelling for t in itertools.takewhile(lambda t:t.spelling!=')', it) if t.kind == ToK.IDENTIFIER]
+              _args = [nm(t) for t in itertools.takewhile(lambda t:nm(t)!=')', it) if clang.clang_getTokenKind(t) == clang.CXToken_Identifier]
               if len(body:=list(it)) == 0: continue
-              macros += [f"{c.spelling} = lambda {','.join(_args)}: {readext(f, body[0].location.offset, toks[-1].extent.end.offset)}"]
-            else: macros += [f"{c.spelling} = {readext(f, toks[1].location.offset, toks[-1].extent.end.offset)}"]
-          case CK.VAR_DECL if c.linkage == LK.INTERNAL:
-            if (c.type.kind == TK.CONSTANTARRAY and c.type.get_array_element_type().get_canonical().kind in ints and
-                (init:=last(c)).kind == CK.INIT_LIST_EXPR and all(re.match(r"\[.*\].*=", readext(f, c.extent)) for c in init.get_children())):
-              cs = init.get_children()
-              macros += [f"{c.spelling} = {{{','.join(f'{readext(f,next(it:=c.get_children()).extent)}:{readext(f,next(it).extent)}' for c in cs)}}}"]
-            elif c.type.get_canonical().kind in ints: macros += [f"{c.spelling} = {readext(f, last(c).extent)}"]
-            else: macros += [f"{c.spelling} = {tname(c.type)}({readext(f, last(c).extent)})"]
-          case CK.VAR_DECL if c.linkage == LK.EXTERNAL and dll:
-            lines.append(f"try: {c.spelling} = {tname(c.type)}.in_dll(dll, '{c.spelling}')\nexcept (ValueError,AttributeError): pass")
+              macros += [f"{nm(c)} = lambda {','.join(_args)}: {readext(f, loc(body[0]), clang.clang_getRangeEnd(extent(toks[-1])))}"]
+            else: macros += [f"{nm(c)} = {readext(f, loc(toks[1]), clang.clang_getRangeEnd(extent(toks[-1])))}"]
+          case clang.CXCursor_VarDecl if clang.clang_getCursorLinkage(c) == clang.CXLinkage_Internal:
+            ty = clang.clang_getCursorType(c)
+            if (ty.kind == clang.CXType_ConstantArray and clang.clang_getCanonicalType(clang.clang_getArrayElementType(ty)).kind in ints and
+                (init:=children(c)[-1]).kind == clang.CXCursor_InitListExpr
+                and all(re.match(r"\[.*\].*=", readext(f, extent(c))) for c in children(init))):
+              cs = children(init)
+              macros += [f"{nm(c)} = {{{','.join(f'{readext(f, extent(next(it:=iter(children(c)))))}:{readext(f, extent(next(it)))}' for c in cs)}}}"]
+            elif clang.clang_getCanonicalType(ty).kind in ints: macros += [f"{nm(c)} = {readext(f, extent(children(c)[-1]))}"]
+            else: macros += [f"{nm(c)} = {tname(ty)}({readext(f, extent(children(c)[-1]))})"]
+          case clang.CXCursor_VarDecl if clang.clang_getCursorLinkage(c) == clang.CXLinkage_External and dll:
+            lines.append(f"try: {nm(c)} = {tname(clang.clang_getCursorType(c))}.in_dll(dll, '{nm(c)}')\nexcept (ValueError,AttributeError): pass")
+          case clang.CXCursor_ObjCProtocolDecl: proto(c)
       except NotImplementedError as e:
-        print(f"skipping {c.spelling}: {e}")
+        print(f"skipping {nm(c)}: {e}")
         lines, types = rollback
+    clang.clang_disposeTranslationUnit(tu)
+    clang.clang_disposeIndex(idx)
   main = (f"# mypy: ignore-errors\nimport ctypes{', os' if any('os' in s for s in dll) else ''}\n"
     "from tinygrad.helpers import unwrap\nfrom tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR\n" + '\n'.join([*prolog,
-      *(["from ctypes.util import find_library"]*any('find_library' in s for s in dll)),
+      *(["from ctypes.util import find_library"]*any('find_library' in s for s in dll)), *(["from tinygrad.runtime.support import objc"]*objc),
       *(["def dll():",*flatten([[f"  try: return ctypes.CDLL(unwrap({d}){', use_errno=True' if use_errno else ''})",'  except: pass'] for d in dll]),
          "  return None", "dll = dll()\n"]*bool(dll)), *lines]) + '\n')
   macros = [r for m in macros if (r:=functools.reduce(lambda s,r:re.sub(r[0], r[1], s), rules + base_rules, m))]

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)
@@ -29,6 +32,7 @@ def CEnum(typ: type[ctypes._SimpleCData]):
 
     def __eq__(self, other): return self.value == other
     def __repr__(self): return self.get(self) if self.value in self.__class__._val_to_name_ else str(self.value)
+    def __hash__(self): return hash(self.value)
 
   return _CEnum
 
@@ -50,16 +54,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 +87,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