fix that test

* generic range changes that work for str + int

* opt range counts up

.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,21 +638,19 @@ 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
-    - name: benchmark MobileNetV2 on DSP
-      run: |
-        # generate quantized weights
-        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
-        ln -s /data/home/tiny/tinygrad/testsig-*.so .
-        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
-        # benchmark on DSP with NOOPT=1, the devectorizer has issues
-        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.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
+    #     ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
+    #     ln -s /data/home/tiny/tinygrad/testsig-*.so .
+    #     PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
+    #     # benchmark on DSP with NOOPT=1, the devectorizer has issues
+    #     PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 

.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:
@@ -306,6 +306,7 @@ jobs:
       with:
         key: spec-unit
         deps: testing_unit
+        python-version: '3.14'
     - name: Test SPEC=2
       run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
 

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

examples/llama3.py

@@ -1,8 +1,6 @@
 from pathlib import Path
 from typing import List
 import json, argparse, random, time, os
-import tiktoken
-from tiktoken.load import load_tiktoken_bpe
 from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
 from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters, gguf_load
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
@@ -12,6 +10,8 @@ from extra.bench_log import BenchEvent, WallTimeEvent
 class Tokenizer:
   pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
   def __init__(self, model_path: str):
+    import tiktoken
+    from tiktoken.load import load_tiktoken_bpe
     mergeable_ranks = load_tiktoken_bpe(model_path)
     self.num_base_tokens = len(mergeable_ranks)
     special_tokens = [

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/amd_uop_matmul.py

@@ -4,9 +4,9 @@ from tinygrad.engine.realize import ExecItem, get_runner
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv
 
-N = 4096
+N = getenv("N", 4096)
 M = K = N
-run_count = 5
+run_count = getenv("CNT", 5)
 
 # ---------------------------
 # launch/config constants
@@ -155,14 +155,15 @@ def test_matmul(sink:UOp, N=N):
       ets.append(ei.run(wait=True))
   print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
 
-  GlobalCounters.reset()
-  with Context(DEBUG=2):
-    tc = (a @ b).realize()
-  with Context(DEBUG=0):
-    err = (hc - tc).square().mean().item()
-  print(f"mean squared error {err}")
-  if err > 1e-06:
-    raise RuntimeError("matmul is wrong!")
+  if getenv("VERIFY", 1):
+    GlobalCounters.reset()
+    with Context(DEBUG=2):
+      tc = (a @ b).realize()
+    with Context(DEBUG=0):
+      err = (hc - tc).square().mean().item()
+    print(f"mean squared error {err}")
+    if err > 1e-06:
+      raise RuntimeError("matmul is wrong!")
 
 if __name__ == "__main__":
   test_matmul(hand_spec_kernel3(), N=N)

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

@@ -1,8 +1,10 @@
-import pathlib
+import os, pathlib
+
+# TODO: there is a timing bug without this
+os.environ["AMD_AQL"] = "1"
+
 from tinygrad.device import Device
 from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
-import time
-import os
 
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
@@ -32,7 +34,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")
 
@@ -44,9 +46,9 @@ if __name__=="__main__":
     raise RuntimeError("Error while initiating AMD device")
 
   COMPILER = HIPCompiler(DEV.arch)
-  if DEV.arch in {'gfx1100', 'gfx1103'}:
-    if DEV.arch == 'gfx1103':
-      NUM_WORKGROUPS = 8
+  if DEV.arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
+    if DEV.arch == 'gfx1103': NUM_WORKGROUPS = 8
+    if DEV.arch == 'gfx1151': NUM_WORKGROUPS = 40
     launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
     launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
     launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))

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,20 +58,23 @@ 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()
-  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)
-  return {int(num, base=16):name for name, num in matches}
+  res = {}
+  known_classes = {"NV01_DEVICE_0", "NV01_ROOT", "NV1_MEMORY_SYSTEM", "NV01_MEMORY_VIRTUAL", "NV1_MEMORY_USER", "NV50_MEMORY_VIRTUAL", "NV_FERMI_VASPACE_A",
+                   "NV20_SUBDEVICE_0"}
+  for nm,val in nv_gpu.__dict__.items():
+    if not isinstance(val, int): continue
+    if 0x3000 < val < 0xffff: res[val] = nm
+    if nm in known_classes: res[val] = nm
+  return res
 nvclasses = get_classes()
 nvuvms = {getattr(nv_gpu, x):x for x in dir(nv_gpu) if x.startswith("UVM_") and nv_gpu.__dict__.get(x+"_PARAMS")}
-nvqcmds = {int(getattr(nv_gpu, x)):x for x in dir(nv_gpu) if x[:7] in {"NVC6C0_", "NVC56F_", "NVC6B5_"} and isinstance(getattr(nv_gpu, x), int)}
+nvqcmds = {int(getattr(nv_gpu, x)):x for x in dir(nv_gpu) if x[:7] in {"NVC9B0_", "NVC6C0_", "NVC56F_", "NVC6B5_"} and isinstance(getattr(nv_gpu, x), int)}
 
 global_ioctl_id = 0
 gpus_user_modes = []
@@ -272,4 +275,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/active_sqtt_parse.py

@@ -0,0 +1,131 @@
+import os
+os.environ["PYTHONPATH"] = "."
+os.environ["SQTT"] = "1"
+if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
+os.environ["PROFILE"] = "1"
+os.environ["AMD_LLVM"] = "0"
+
+from dataclasses import replace
+import atexit, contextlib
+from tinygrad import Tensor
+from tinygrad.helpers import system, OSX
+from tinygrad.runtime.ops_amd import AMDProgram
+from extra.sqtt.roc import decode, WaveExec, ProfileSQTTEvent
+from tinygrad.device import Device, ProfileDeviceEvent
+
+from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
+
+# TODO: should really check for AM driver / USB
+if not OSX:
+  def set_power(x): system(f"sudo /opt/rocm/bin/amd-smi set -l {x}")
+  @atexit.register
+  def reset_power(): set_power("auto")
+  set_power("stable_std")
+
+dev = Device["AMD"]
+
+@contextlib.contextmanager
+def save_sqtt():
+  # clear the old traces
+  dev.profile_events.clear()
+  sqtt:dict[str, list[WaveExec]] = {}
+  yield sqtt
+  events = dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())]
+
+  rctx = decode(events)
+  assert len(rctx.inst_execs) > 0, "empty sqtt output"
+  sqtt.update(rctx.inst_execs)
+
+  for e in events:
+    if isinstance(e, ProfileSQTTEvent):
+      print(replace(e, blob=b''))
+      if e.se == 0:
+        parse_sqtt_print_packets(e.blob)
+
+template = """.text
+.globl matmul
+.p2align 8
+.type matmul,@function
+matmul:
+  INSTRUCTION
+  s_endpgm
+
+.rodata
+.p2align 6
+.amdhsa_kernel matmul
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
+  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
+  .amdhsa_wavefront_size32 1
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.version:
+  - 1
+  - 0
+amdhsa.kernels:
+  - .name: matmul
+    .symbol: matmul.kd
+    .group_segment_fixed_size: 0
+    .private_segment_fixed_size: 0
+    .wavefront_size: 32
+    .sgpr_count: 8
+    .vgpr_count: 32
+    .max_flat_workgroup_size: 1024
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 8
+    .args:
+      - .address_space:  global
+        .name:           a
+        .offset:         0
+        .size:           8
+        .type_name:      'float*'
+        .value_kind:     global_buffer
+...
+.end_amdgpu_metadata
+"""
+
+def run_asm(src):
+  NUM_WORKGROUPS = 1
+  WAVE_SIZE = 32
+  NUM_WAVES = 1
+  t = Tensor.empty(0x1000).realize()
+  buf = t.uop.buffer.ensure_allocated()
+  lib = dev.compiler.compile(template.replace("INSTRUCTION", '\n'.join(src)))
+  dev.compiler.disassemble(lib)
+  fxn = AMDProgram(dev, "matmul", lib)
+  fxn(buf._buf, global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True)
+
+if __name__ == "__main__":
+  with save_sqtt() as sqtt:
+    #(Tensor.empty(16,16) @ Tensor.empty(16,16)).elu().realize()
+    Tensor.empty(1).elu().realize()
+  exit(0)
+
+  with save_sqtt() as sqtt:
+    # what's in v0?
+    run_asm([
+      "v_mov_b32_e32 v0, 0",
+      "v_mov_b32_e32 v1, 0",
+      "s_clause 0x1",
+      "s_load_b64 s[0:1], s[0:1], null",
+      "s_waitcnt lgkmcnt(0)",
+    ]+[
+      "global_load_b32 v1, v0, s[0:1]",
+    ]*10+[
+      "global_load_b32 v10, v1, s[0:1]",
+      "s_waitcnt vmcnt(0)",
+
+      #"v_rcp_f32 v1, v0"
+      #"v_add_f32_e32 v1 v0 v0",
+      #"v_add_f32_e32 v5 v4 v4",
+      #"v_add_f32_e32 v7 v6 v6",
+      #"v_add_f32_e32 v1 v0 v0",
+      #"v_add_f32_e32 v2 v1 v1",
+      #"s_nop 1"
+    ]*5+[
+      "v_add_f32_e32 v3 v2 v2",
+    ]*5+[
+      "v_mul_f32_e32 v3 v2 v2",
+    ]*7)

extra/sqtt/attempt_sqtt_parse.py

@@ -0,0 +1,543 @@
+import pickle
+from tinygrad.helpers import getenv
+from extra.sqtt.roc import decode, ProfileSQTTEvent
+
+# Instruction packets (one per ISA op)
+# NOTE: these are bad guesses and may be wrong! feel free to update if you know better
+# some names were taken from SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT
+
+OPCODE_NAMES = {
+  # gated by SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT
+  0x02: "VMEMEXEC",
+  # gated by SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT
+  0x03: "ALUEXEC",
+  # gated by SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT (but others must be enabled for it to show)
+  0x01: "VALUINST",
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT
+  0x06: "WAVERDY",
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT
+  0x08: "WAVEEND",
+  0x09: "WAVESTART",
+  # gated by SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT
+  0x04: "IMMEDIATE_4",
+  0x05: "IMMEDIATE_5",
+  # some gated by SQ_TT_TOKEN_EXCLUDE_REG_SHIFT, some always there
+  0x14: "REG",
+  # gated by SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT
+  0x12: "EVENT",
+  # gated by SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
+  0x18: "INST",
+  # gated by SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT
+  0x19: "UTILCTR",
+
+  # ------------------------------------------------------------------------
+  # 0x07–0x0F: pure timestamp-ish deltas
+  # ------------------------------------------------------------------------
+  0x07: "TS_DELTA_S8_W3",           # shift=8,  width=3  (small delta)
+  0x0A: "TS_DELTA_S5_W2_A",         # shift=5,  width=2
+  0x0B: "TS_DELTA_S5_W3_A",         # shift=5,  width=3
+  0x0C: "TS_DELTA_S5_W3_B",         # shift=5,  width=3 (different consumer)
+  0x0D: "TS_DELTA_S5_W3_C",         # shift=5,  width=3
+  0x0E: "TS_DELTA_S7_W2",           # shift=7,  width=2
+  0x0F: "TS_DELTA_SHORT_PLUS4",     # short delta; ROCm adds +4 before accumulate
+
+  # ------------------------------------------------------------------------
+  # 0x10–0x19: timestamps, layout headers, events, perf
+  # ------------------------------------------------------------------------
+  0x10: "PSEUDO_NEED_MORE_BITS",    # not a real packet; decoder refill hint
+
+  0x11: "TS_WAVE_STATE_SAMPLE",     # wave stall/termination sample (byte at +10)
+  0x13: "EVT_SMALL_GENERIC",        # same structural family as 0x08/0x12/0x19
+
+  0x15: "PERFCOUNTER_SNAPSHOT",     # small delta + 50-ish bits of snapshot
+  0x16: "TS_DELTA36_OR_MARK",       # 36-bit long delta or 36-bit marker
+  0x17: "LAYOUT_MODE_HEADER",       # layout/mode/group + selectors A/B
+}
+
+# these tables are from rocprof trace decoder
+# rocprof_trace_decoder_parse_data-0x11c6a0
+# parse_sqtt_180 = b *rocprof_trace_decoder_parse_data-0x11c6a0+0x110040
+
+# ---------- 1. local_138: 256-byte state->token table ----------
+
+STATE_TO_TOKEN: bytes = bytes([
+  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x12, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x13, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+])
+
+
+# ---------- 2. DAT_0012e280: nibble budget per opcode&0x1F ----------
+
+NIBBLE_BUDGET = [
+  0x08, 0x0C, 0x08, 0x08, 0x0C, 0x18, 0x18, 0x40,
+  0x14, 0x20, 0x30, 0x14, 0x34, 0x1C, 0x30, 0x08,
+  0x04, 0x18, 0x18, 0x20, 0x40, 0x40, 0x30, 0x40,
+  0x14, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+]
+assert len(NIBBLE_BUDGET) == 32
+
+
+# ---------- 3. delta_map from your hash nodes ----------
+
+# opcode -> (shift, width)
+DELTA_MAP_DEFAULT = {
+  0x01: (3,  3),   # shift=3,  end=6
+  0x02: (4,  2),   # shift=4,  end=6
+  0x03: (4,  2),   # shift=4,  end=6
+  0x04: (4,  3),   # shift=4,  end=7
+  0x05: (5,  3),   # shift=5,  end=8
+  0x06: (5,  3),   # shift=5,  end=8
+  0x07: (8,  3),   # shift=8,  end=11
+  0x08: (5,  3),   # shift=5,  end=8
+  0x09: (5,  2),   # shift=5,  end=7
+  0x0A: (5,  2),   # shift=5,  end=7
+  0x0B: (5,  3),   # shift=5,  end=8
+  0x0C: (5,  3),   # shift=5,  end=8
+  0x0D: (5,  3),   # shift=5,  end=8
+  0x0E: (7,  2),   # shift=7,  end=9
+  0x0F: (4,  4),   # shift=4,  end=8
+  0x10: (0,  0),   # shift=0,  end=0  (no delta)
+  0x11: (7,  9),   # shift=7,  end=16
+  0x12: (8,  3),   # shift=8,  end=11
+  0x13: (8,  3),   # shift=8,  end=11
+  0x14: (4,  3),   # shift=4,  end=7
+  0x15: (7,  3),   # shift=7,  end=10
+  0x16: (12, 36),  # shift=12, end=48 (36-bit field, matches the 0x16 special-case)
+  0x17: (0,  0),   # shift=0,  end=0  (no delta)
+  0x18: (4,  3),   # shift=4,  end=7
+  0x19: (7,  2),   # shift=7,  end=9
+}
+
+# ---------- 4. One-line-per-packet parser ----------
+
+def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
+  """
+  Decode packet payloads conservatively, using:
+    - NIBBLE_BUDGET[opcode & 0x1F] to mask reg down to true width.
+    - DELTA_MAP_DEFAULT[opcode] to expose the "primary" field (often delta).
+    - Per-opcode layouts derived from rocprof's decompiled consumers.
+  """
+  # --- 0. Restrict to real packet bits ---------------------------------
+  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]
+  if nb_bits <= 0 or nb_bits >= 64:
+    pkt = reg & ((1 << 64) - 1)
+  else:
+    pkt = reg & ((1 << nb_bits) - 1)
+
+  fields: list[str] = []
+
+  shift, width = DELTA_MAP_DEFAULT.get(opcode, (0, 0))
+  if width:
+    field_mask = (1 << width) - 1
+    shaped_field = (pkt >> shift) & field_mask
+  else:
+    field_mask = 0
+    shaped_field = 0
+
+  # =====================================================================
+  # 1. Timestamp-centric opcodes (actually drive 'time')
+  # =====================================================================
+
+  if opcode == 0x0F:  # TS_DELTA_SHORT_PLUS4
+    # In the caller, delta already has +4 applied.
+    raw_delta = shaped_field
+    fields.append(f"raw_delta={raw_delta}")
+    fields.append(f"ts_short_plus4={delta}")
+    return ", ".join(fields)
+
+  if opcode == 0x11:  # TS_WAVE_STATE_SAMPLE
+    # DELTA_MAP_DEFAULT: shift=7, width=9 -> small delta.
+    raw_delta = shaped_field
+    coarse    = (pkt >> (shift + width)) & 0xFF  # matches byte at +10 in C
+    fields.append(f"raw_delta={raw_delta}")
+    if coarse:
+      fields.append(f"coarse_state=0x{coarse:02x}")
+    # From decomp:
+    #  - when layout<3 and coarse&1, it sets a "has interesting wave" flag
+    #  - when coarse&8, it marks all live waves as "terminated"
+    if coarse & 0x01:
+      fields.append("flag_wave_interest=1")
+    if coarse & 0x08:
+      fields.append("flag_terminate_all=1")
+    return ", ".join(fields)
+
+  if opcode == 0x16:  # TS_DELTA36_OR_MARK
+    # Bits:
+    #   bit8  -> 0x100
+    #   bit9  -> 0x200
+    #   bits 12..47 -> 36-bit field used as delta or marker
+    bit8 = bool(pkt & 0x100)
+    bit9 = bool(pkt & 0x200)
+    if not bit9:
+      mode = "delta"
+    elif not bit8:
+      mode = "marker"
+    else:
+      mode = "other"
+    val36 = (pkt >> 12) & ((1 << 36) - 1)
+    fields.append(f"mode={mode}")
+    if mode != "delta":
+      fields.append(f"val36=0x{val36:x}")
+    return ", ".join(fields)
+
+  # For 0x07, 0x0A–0x0E, we know they drive time (via DELTA_MAP_DEFAULT),
+  # but we don't see any other fields used in the decomp.
+  if opcode in (0x07, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E):
+    if width:
+      raw_delta = shaped_field
+      leftover  = pkt & ~(field_mask << shift)
+      fields.append(f"raw_delta={raw_delta}")
+      if leftover:
+        fields.append(f"payload=0x{leftover:x}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 2. Small "meta + tiny delta" packets (0x01–0x06)
+  # =====================================================================
+
+  if opcode == 0x01:  # META_ID12_TS_SMALL
+    id12 = pkt & 0xFFF
+    fields.append(f"id12=0x{id12:03x}")
+    if width:
+      fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    return ", ".join(fields)
+
+  if opcode == 0x02:  # META_FLAG8_TS_SMALL
+    flag8 = pkt & 0xFF
+    fields.append(f"flag8=0x{flag8:02x}")
+    if width:
+      fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    return ", ".join(fields)
+
+  if opcode == 0x03:  # META_SUBEVENT8_TS_SMALL
+    sub8 = pkt & 0xFF
+    fields.append(f"subevent8=0x{sub8:02x}")
+    if width:
+      fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    return ", ".join(fields)
+
+  if opcode == 0x04:  # META_BASE_INDEX12_TS
+    idx12 = pkt & 0xFFF
+    fields.append(f"base_index12=0x{idx12:03x}")
+    if width:
+      fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    return ", ".join(fields)
+
+  if opcode in (0x05, 0x06):  # META_DESC24_TS_A/B
+    desc24 = pkt & 0xFFFFFF
+    fields.append(f"desc24=0x{desc24:06x}")
+    if width:
+      fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 3. Opcode 0x14: exec/config record (+ COR marker)
+  # =====================================================================
+
+  if opcode == 0x14:  # INST_EXEC_OR_CFG
+    subop   = (pkt >> 16) & 0xFFFF       # (short)(w >> 0x10)
+    val32   = (pkt >> 32) & 0xFFFFFFFF   # (uint)(w >> 0x20)
+    slot    = (pkt >> 7) & 0x7           # index in local_168[...] tables
+    hi_byte = (pkt >> 8) & 0xFF          # determines config vs marker
+
+    fields.append(f"subop=0x{subop:04x}")
+    fields.append(f"slot={slot}")
+    fields.append(f"val32=0x{val32:08x}")
+
+    if hi_byte & 0x80:
+      # Config flavour: writes config words into per-slot state arrays.
+      fields.append("kind=config")
+      if subop == 0x000C:
+        fields.append("cfg_target=local_168[slot].lo")
+      elif subop == 0x000D:
+        fields.append("cfg_target=local_168[slot].hi")
+    else:
+      # COR marker: subop 0xC342, payload "COR\0" → start of a COR region.
+      if subop == 0xC342:
+        fields.append("kind=cor_stream")
+        if val32 == 0x434F5200:
+          fields.append("cor_magic='COR\\0'")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 4. Opcode 0x17: layout / mode header
+  # =====================================================================
+
+  if opcode == 0x17:  # LAYOUT_MODE_HEADER
+    # From decomp (two sites with identical logic):
+    #   layout = (w >> 7) & 0x3f
+    #   mode   = (w >> 0xd) & 3
+    #   group  = (w >> 0xf) & 7
+    #   sel_a  = (w >> 0x1c) & 0xf
+    #   sel_b  = (w >> 0x21) & 7
+    #   flag4  = (w >> 0x3b) & 1  (only meaningful when layout == 4)
+    layout = (pkt >> 7)  & 0x3F
+    mode   = (pkt >> 13) & 0x3
+    group  = (pkt >> 15) & 0x7
+    sel_a  = (pkt >> 0x1C) & 0xF
+    sel_b  = (pkt >> 0x21) & 0x7
+    flag4  = (pkt >> 0x3B) & 0x1
+
+    fields.append(f"layout={layout}")
+    fields.append(f"group={group}")
+    fields.append(f"mode={mode}")
+    fields.append(f"sel_a={sel_a}")
+    fields.append(f"sel_b={sel_b}")
+    if layout == 4:
+      fields.append(f"layout4_flag={flag4}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 5. Opcode 0x09: state / route config record
+  # =====================================================================
+
+  if opcode == 0x09:  # PERF_ROUTE_CONFIG
+    # From case 9 in multiple consumers:
+    #   flag7  = (w >> 7) & 1        (low bit of uVar41)
+    #   cls2   = (w >> 8) & 3        (class / group)
+    #   slot4  = (w >> 10) & 0xf     (slot / group index)
+    #   idx_lo = (w >> 0xd) & 0x1f   (low index, layout<4 path)
+    #   idx_hi = (w >> 0xf) & 0x1f   (high index, layout>=4 path)
+    #   id7    = (w >> 0x19) & 0x7f  (7-bit id)
+    flag7   = (pkt >> 7) & 0x1
+    cls2    = (pkt >> 8) & 0x3
+    slot4   = (pkt >> 10) & 0xF
+    idx_lo  = (pkt >> 13) & 0x1F
+    idx_hi  = (pkt >> 15) & 0x1F
+    id7     = (pkt >> 0x19) & 0x7F
+
+    fields.append(f"flag7={flag7}")
+    fields.append(f"cls2={cls2}")
+    fields.append(f"slot4=0x{slot4:x}")
+    fields.append(f"idx_lo5=0x{idx_lo:x}")
+    fields.append(f"idx_hi5=0x{idx_hi:x}")
+    fields.append(f"id7=0x{id7:x}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 6. Opcode 0x18: perf/event selector (FUN_0010aba0)
+  # =====================================================================
+
+  if opcode == 0x18:  # PERF_EVENT_SELECT
+    # From case 0x18:
+    #   low3   = w & 7
+    #   grp3   = (w >> 3) or (w >> 4) & 7   (layout-dependent)
+    #   flags  = bits 6 (B6) and 7 (B7)
+    #   hi8    = (w >> 0xc) & 0xff   (layout 4 path)
+    #   hi7    = (w >> 0xd) & 0x7f   (other layouts)
+    #   idx5   = (w >> 7) or (w >> 8) & 0x1f, used as wave index
+    low3     = pkt & 0x7
+    grp3_a   = (pkt >> 3) & 0x7
+    grp3_b   = (pkt >> 4) & 0x7
+    flag_b6  = (pkt >> 6) & 0x1
+    flag_b7  = (pkt >> 7) & 0x1
+    idx5_a   = (pkt >> 7) & 0x1F
+    idx5_b   = (pkt >> 8) & 0x1F
+    hi8      = (pkt >> 12) & 0xFF
+    hi7      = (pkt >> 13) & 0x7F
+
+    fields.append(f"low3=0x{low3:x}")
+    fields.append(f"grp3_a=0x{grp3_a:x}")
+    fields.append(f"grp3_b=0x{grp3_b:x}")
+    fields.append(f"flag_b6={flag_b6}")
+    fields.append(f"flag_b7={flag_b7}")
+    fields.append(f"idx5_a=0x{idx5_a:x}")
+    fields.append(f"idx5_b=0x{idx5_b:x}")
+    fields.append(f"hi8=0x{hi8:02x}")
+    fields.append(f"hi7=0x{hi7:02x}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 7. Opcode 0x15: perfcounter snapshot
+  # =====================================================================
+
+  if opcode == 0x15:  # PERFCOUNTER_SNAPSHOT
+    # NIBBLE_BUDGET gives full 64 bits here.
+    # DELTA_MAP_DEFAULT: shift=7, width=3 → tiny delta field.
+    raw_delta = shaped_field if width else 0
+    # low bits below the delta field
+    snap_low  = pkt & ((1 << shift) - 1) if shift else 0
+    # everything above delta field
+    snap_hi   = pkt >> (shift + width) if width else (pkt >> shift)
+
+    fields.append(f"raw_delta={raw_delta}")
+    fields.append(f"snap_low_s{shift}=0x{snap_low:x}")
+    fields.append(f"snap_hi=0x{snap_hi:x}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 8. Small event-ish packets (0x08 / 0x12 / 0x13 / 0x19)
+  # =====================================================================
+
+  if opcode in (0x08, 0x12, 0x13, 0x19):
+    # These are all "small event / metric" style tokens. The exact semantics
+    # depend on layout (0x17) and accumulated state (local_500 etc), so we
+    # expose:
+    #   - low 8 bits as kind byte
+    #   - rest as opaque payload.
+    kind    = pkt & 0xFF
+    payload = pkt >> 8
+    fields.append(f"kind_byte=0x{kind:02x}")
+    if payload:
+      fields.append(f"payload=0x{payload:x}")
+    return ", ".join(fields)
+
+  # =====================================================================
+  # 9. Pseudo opcode 0x10: never a "real" packet
+  # =====================================================================
+
+  if opcode == 0x10:  # PSEUDO_NEED_MORE_BITS
+    # The main loop never prints these; they're just a control token.
+    return ""
+
+  # =====================================================================
+  # 10. Generic fallback: expose the DELTA_MAP_DEFAULT field + leftover
+  # =====================================================================
+
+  if width:
+    fields.append(f"field_s{shift}_w{width}={shaped_field}")
+    leftover = pkt & ~(field_mask << shift)
+    if leftover:
+      fields.append(f"payload=0x{leftover:x}")
+
+  return ", ".join(fields)
+
+# 0xb is time something
+# 0xd is time something
+# 0xf is small time advance
+# 0x11 is time advance
+# 0x16 is big time advance + markers
+# 0x14 is REG
+DEFAULT_FILTER = (0xb, 0xd, 0xf, 0x11, 0x16, 0x14) if getenv("FILTER", 1) else None
+
+def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000, filter=DEFAULT_FILTER) -> None:
+  """
+  Minimal debug: print ONE LINE per decoded token (packet).
+
+  Now prints only the actual nibbles that belong to each packet, instead of
+  the full 64-bit shift register.
+  """
+  n = len(data)
+  time = 0
+  reg = 0          # shift register
+  offset = 0       # bit offset, in steps of 4 (one nibble)
+  nib_budget = 0x40
+  flags = 0
+  token_index = 0
+
+  while (offset >> 3) < n and token_index < max_tokens:
+    # Remember where we started refilling for this step (bit offset),
+    # but the *logical* start of the current packet is last_real_offset.
+    refill_start = offset
+
+    # 1) Fill register with nibbles according to nib_budget
+    if nib_budget != 0:
+      target = refill_start + 4 + ((nib_budget - 1) & ~3)
+      cur = refill_start
+      while cur != target and (cur >> 3) < n:
+        byte_index = cur >> 3
+        byte = data[byte_index]
+        shift = 4 if (cur & 4) else 0  # low then high nibble
+        nib = (byte >> shift) & 0xF
+        reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
+        cur += 4
+      offset = cur
+
+    # 2) Decode token from low 8 bits
+    state = reg & 0xFF
+    opcode = STATE_TO_TOKEN[state]
+
+    # 3) Handle pseudo-token 0x10: need more bits, don't print. Looks like a NOP.
+    if opcode == 0x10:
+      # "need more bits" pseudo-token: adjust nibble budget and continue
+      nib_budget = 4
+      if (offset >> 3) >= n:
+        break
+      # Do NOT count this as a real packet; do not update last_real_offset.
+      continue
+
+    # 4) Set next nibble budget
+    nb_index = opcode & 0x1F
+    nib_budget = NIBBLE_BUDGET[nb_index]
+    time_before = time
+    note = ""
+    # 5) Special opcode 0x16 (timestamp / marker)
+    if opcode == 0x16:
+      two_bits = (reg >> 8) & 0x3
+      if two_bits == 1:
+        flags |= 0x01
+
+      # Common 36-bit field at bits [12..47]
+
+      if (reg & 0x200) == 0:
+        # delta mode: add 36-bit delta to time
+        delta = (reg >> 12) & ((1 << 36) - 1)
+        time += delta
+      else:
+        # marker / other modes: no time advance
+        if (reg & 0x100) == 0:
+          # real marker: bit9=1, bit8=0, non-zero payload
+          # "other" 0x16 variants, ignored for timing
+          delta = 0
+    else:
+      # 6) Generic opcode (including 0x0F)
+      shift, width = DELTA_MAP_DEFAULT[opcode]
+      mask = (1 << width) - 1
+      delta = (reg >> shift) & mask
+
+      # TODO: add more opcode parsers here that add notes to other opcodes
+      if opcode == 0x0F:
+        delta_with_fix = delta + 4
+        time += delta_with_fix
+        delta = delta_with_fix
+      else:
+        time += delta
+
+    # Append extra decoded fields into the note string
+    note = decode_packet_fields(opcode, reg, delta)
+
+    if filter is None or opcode not in filter:
+      my_reg = reg
+      my_reg &= (1 << nib_budget) - 1
+      print(
+        f"{token_index:4d}  "
+        f"off={offset//4:5d}  "
+        f"op=0x{opcode:02x} "
+        f"{OPCODE_NAMES[opcode]:24s} "
+        f" time={time_before:8d}+{delta:8d}  "
+        f"{my_reg:16X} "
+        f"{note}"
+      )
+
+    token_index += 1
+
+  # Optional summary at the end
+  print(f"# done: tokens={token_index}, final_time={time}, flags=0x{flags:02x}")
+
+def parse(fn:str):
+  dat = pickle.load(open(fn, "rb"))
+  ctx = decode(dat)
+  dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
+  print(f"got {len(dat_sqtt)} SQTT events in {fn}")
+  return dat_sqtt
+
+if __name__ == "__main__":
+  #dat_sqtt = parse("extra/sqtt/examples/profile_empty_run_0.pkl")
+  #dat_sqtt = parse("extra/sqtt/examples/profile_plus_run_0.pkl")
+  dat_sqtt = parse("extra/sqtt/examples/profile_gemm_run_0.pkl")
+  blob_0 = dat_sqtt[0].blob
+  parse_sqtt_print_packets(blob_0[8:])

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

@@ -1,4 +1,4 @@
-import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools
+import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools, threading
 from tinygrad.helpers import temp, unwrap, DEBUG
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
@@ -41,6 +41,7 @@ class WaveExec:
   wave_id:int
   cu:int
   simd:int
+  se:int
   begin_time:int
   end_time:int
   insts:list[InstExec]
@@ -63,21 +64,23 @@ class _ROCParseCtx:
     self.active_blob = (ctypes.c_ubyte * len(x.blob)).from_buffer_copy(x.blob) if x is not None else None
     return self.active_blob
 
-  def on_occupancy_ev(self, ev):
+  def on_occupancy_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_occupancy_t):
     if DEBUG >= 5: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
 
-  def on_wave_ev(self, ev):
+  def on_wave_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_wave_t):
     if DEBUG >= 5: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
 
     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))
+      if DEBUG >= 8: print(inst_execs[-1])
 
     if ev.instructions_size > 0:
-      self.inst_execs.setdefault(unwrap(self.active_kern), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, ev.begin_time, ev.end_time, inst_execs))
+      self.inst_execs.setdefault(unwrap(self.active_kern), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
+                                                                               ev.end_time, inst_execs))
 
 def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
   dev_events:dict[str, ProfileDeviceEvent] = {}
@@ -91,25 +94,25 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
   ROCParseCtx = _ROCParseCtx(dev_events, sqtt_events, prog_events)
 
   @rocprof.rocprof_trace_decoder_se_data_callback_t
-  def copy_cb(buf, buf_size, data_ptr):
+  def copy_cb(buf, buf_size, _):
     if (prof_info:=ROCParseCtx.next_sqtt()) is None: return 0
     buf[0] = ctypes.cast(prof_info, ctypes.POINTER(ctypes.c_ubyte))
     buf_size[0] = len(prof_info)
     return len(prof_info)
 
   @rocprof.rocprof_trace_decoder_trace_callback_t
-  def trace_cb(record_type, events_ptr, n, data_ptr):
+  def trace_cb(record_type, events_ptr, n, _):
     match record_type:
       case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
         for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr): ROCParseCtx.on_occupancy_ev(ev)
       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
-  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
+  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
     instr, mem_size_ptr[0] = ROCParseCtx.disasms[(unwrap(ROCParseCtx.active_kern), pc.address)]
 
     # this is the number of bytes to next instruction, set to 0 for end_pgm
@@ -123,9 +126,11 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
 
     return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
 
-  try:
-    rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-  except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
+  def worker():
+    try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
+    except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
+  (t:=threading.Thread(target=worker, daemon=True)).start()
+  t.join()
   return ROCParseCtx
 
 if __name__ == "__main__":

extra/sqtt/test_timing.py

@@ -7,14 +7,12 @@ os.environ["AMD_LLVM"] = "0"
 
 import unittest
 import sys, contextlib
-from tinygrad import Tensor
-from tinygrad.dtype import dtypes
-from tinygrad.renderer import ProgramSpec
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AddrSpace
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad import Tensor, dtypes
+from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 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,13 +34,13 @@ def asm_kernel(instrs:list[str], l:int=1, g:int=1) -> Tensor:
 def save_sqtt():
   # clear the old traces
   dev.profile_events.clear()
-  sqtt:dict[PrgExec, list[InstExec]] = {}
+  sqtt:dict[str, list[WaveExec]] = {}
   yield sqtt
   # decode sqtt
-  if os.environ["DEV"] == "AMD":
-    rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
-    assert len(rctx.inst_execs) > 0, "empty sqtt output"
-    sqtt.update(rctx.inst_execs)
+  if os.environ["DEV"] != "AMD": return
+  rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
+  assert len(rctx.inst_execs) > 0, "empty sqtt output"
+  sqtt.update(rctx.inst_execs)
 
 class TestTiming(unittest.TestCase):
   def test_v_add(self):
@@ -75,7 +73,6 @@ class TestTiming(unittest.TestCase):
     inp = Tensor([-2.0]).realize()
     with save_sqtt() as sqtt:
       Tensor.custom_kernel(out, inp, fxn=custom_vrcp)[0].realize()
-
     wave = list(sqtt.values())[0][0]
     for i in range(len(wave.insts)):
       if wave.insts[i].inst.startswith("global_store"):
@@ -84,13 +81,17 @@ class TestTiming(unittest.TestCase):
 
   def test_wmma(self):
     with save_sqtt() as sqtt:
-      asm_kernel([
-        "v_wmma_f32_16x16x16_f16 v[16:23], v[0:7], v[8:15], v[16:23]",
-        "v_add_f32_e32 v0 v16 v0",
-      ], l=32*4).realize()
-    assert len(sqtt) == 2, f"expected two waves, got {len(sqtt)} {list(sqtt.keys())}"
-    wmma = list(sqtt.values())[0][0]
-    self.assertGreater(wmma.dur, 1) # rgp says 32 clocks
+      for tc in dev.renderer.get_tensor_cores(dev.arch):
+        M, K, N = tc.dims
+        s = 32
+        a = Tensor.empty(M*s, K*s, dtype=tc.dtype_in)@Tensor.empty(K*s, N*s, dtype=tc.dtype_in)
+        a.realize()
+        print(a)
+    for p,waves in sqtt.items():
+      for e in waves[0].insts:
+        if (e.inst.startswith("v_wmma")):
+          instruction = e.inst.split(" ")[0]
+          print(f"{instruction:<29} : {e.dur} cycles")
 
   def test_sleep(self):
     n = 1
@@ -98,15 +99,35 @@ class TestTiming(unittest.TestCase):
       assert data0.dtype.base == dtypes.ulong
       op = custom("unsigned long long t0 = __builtin_readcyclecounter();")
       op = custom(f"__builtin_amdgcn_s_sleep({n});", op)
-      op = custom(f"unsigned long long t1 = __builtin_readcyclecounter();", op)
+      op = custom("unsigned long long t1 = __builtin_readcyclecounter();", op)
       op = custom(f"data0_{data0.size}[0] = t1 - t0;", op)
       return UOp.sink(data0, op, arg=KernelInfo(name=f"sleep_{n}"))
     diff_hw_reg = Tensor.empty(1, dtype=dtypes.ulong)
     diff_hw_reg = Tensor.custom_kernel(diff_hw_reg, fxn=sleep_kernel)[0]
     with save_sqtt() as sqtt:
       diff_hw_reg.realize()
-    diff_sqtt = list(sqtt.values())[0][2]
-    self.assertEqual(diff_sqtt.dur, diff_hw_reg.item()-1) # 1 cycle for reading the counter register
+    sleep = next((e for e in sqtt[f"sleep_{n}"][0].insts if e.inst.startswith("s_sleep")))
+    # cycles = sleep dur + overhead of storing hi/lo REG_SHADER_CYCLES
+    self.assertGreaterEqual(diff_hw_reg.item(), sleep.dur)
+
+  def test_nop(self):
+    with save_sqtt() as sqtt:
+      asm_kernel(["s_nop 1"]*10).realize()
+    wave = list(sqtt.values())[0][0]
+    for e in wave.insts:
+      print(f"{e.inst} {e.dur=} {e.stall=}")
+
+  def test_wave_sched(self):
+    num_waves = getenv("NUM_WAVES", 16)
+    num_wgps = getenv("NUM_WGPS", 2)
+    num_vgpr = getenv("NUM_VGPR", 256)
+    with save_sqtt() as sqtt:
+      # 1 cycle decode, no stall
+      asm_kernel([f"v_mov_b32_e32 v{i} {i}" for i in range(num_vgpr)], l=32*num_waves, g=num_wgps).realize()
+    waves = list(sqtt.values())[0]
+    print(len(waves), "waves decoded")
+    for w in waves:
+      print(f"{w.wave_id:<2} {w.simd=} {w.cu=} {w.se=} @ clk {w.begin_time}")
 
 if __name__ == "__main__":
   unittest.main()

extra/thunder/tiny/tk/__init__.py

@@ -1 +1,6 @@
-WARP_THREADS = 32
+from tinygrad.device import Device
+
+if Device.DEFAULT == "AMD":
+  WARP_THREADS = 64
+else:
+  WARP_THREADS = 32

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):
+    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
+    return c.after(c_store).reshape(c.shape)
 
-  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):
+    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
+    return c.after(c_store).reshape(c.shape)
 
   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,68 +134,69 @@ 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)
 
   # 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):
+  LOAD_INNER = 4
+  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):
+  STORE_INNER = 4
+  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,11 @@ 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()
+    last_store = self.store_stack.pop()[0]
+    if hasattr(last_store, '_uop'): uop = last_store._uop
+    else: uop = last_store
+
+    return 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/external/external_test_example.py

@@ -1,13 +1,14 @@
 import unittest
 from tinygrad import Device
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv, CI
+from tinygrad.helpers import getenv, CI, OSX
 
 def multidevice_test(fxn):
   exclude_devices = getenv("EXCLUDE_DEVICES", "").split(",")
   def ret(self):
     for device in Device._devices:
-      if device in ["REMOTE", "DISK", "NPY", "FAKE", "DSP", "NULL"]: continue
+      # broken on OSX USB AMD, why?
+      if device in ["REMOTE", "DISK", "NPY", "FAKE", "DSP", "NULL"] or (OSX and device in ["AMD"]): continue
       if not CI: print(device)
       if device in exclude_devices:
         if not CI: print(f"WARNING: {device} test is excluded")

test/external/process_replay/process_replay.py

@@ -36,7 +36,9 @@ def trunc_log(x):
   logging.info("\n".join(lines))
 
 # user config
-SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
+# NOTE: process replay is slow so it's now disabled by default. add [pr] to enable it
+#SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
+SKIP_PROCESS_REPLAY = not ASSERT_DIFF
 if REF == "master": SKIP_PROCESS_REPLAY = True
 class ProcessReplayWarning(Warning): pass
 

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)
@@ -124,6 +124,7 @@ class PM4Executor(AMDQueue):
     elif mem_data_sel == 3:
       if mem_event_type == CACHE_FLUSH_AND_INV_TS_EVENT: ptr.cast('Q')[0] = int(time.perf_counter() * 1e8)
       else: raise RuntimeError(f"Unknown {mem_data_sel=} {mem_event_type=}")
+    elif mem_data_sel == 0: pass # no write
     else: raise RuntimeError(f"Unknown {mem_data_sel=}")
 
   def _exec_copy_data(self, n):

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

@@ -5,26 +5,40 @@ from tinygrad.helpers import CI, Context, getenv
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.uop.ops import Ops
+from tinygrad.renderer import Estimates
+from tinygrad.renderer.ptx import PTXRenderer
 
 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)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexing is weird")
+  def test_tri_complexity(self):
+    with Context(NOOPT=1):
+      t = Tensor.ones(256, 256).contiguous().realize()
+      sched = t.triu().schedule()
+      p = get_program(sched[-1].ast)
+      self.assertLessEqual(Estimates.from_uops(p.uops).ops, 4 * 256 * 256)
+
 DSET, DDIM = 2048, 32
 
 class TestIndexing(unittest.TestCase):

test/test_custom_kernel.py

@@ -102,6 +102,11 @@ def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
 # **** tests ****
 
 class TestCustomKernel(unittest.TestCase):
+  def test_empty(self):
+    a = Tensor.empty(1)
+    a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink())[0]
+    a.realize()
+
   def test_simple(self):
     a = Tensor.ones(16, 16).contiguous()
     b = Tensor.ones(16, 16).contiguous()

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

@@ -1,5 +1,6 @@
 import unittest
-from tinygrad import Tensor, UOp
+import numpy as np
+from tinygrad import Tensor, UOp, nn
 from tinygrad.uop.ops import AxisType, Ops
 
 class TestOuterworldReduce(unittest.TestCase):
@@ -11,6 +12,81 @@ 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))
+    comp = vec_i.contiguous() @ mats[i]
+    store = vec_i.uop.store(comp.uop).end(i)
+    out = Tensor(vec.uop.after(store))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+
+class TestOuterScan(unittest.TestCase):
+  def _test_scan(self):
+    vec = Tensor.randn(1, 10).realize()
+    mats = Tensor.randn(3, 10, 10).realize()
+
+    # 3 matmuls in "scan"
+    vec1 = vec @ mats[0]
+    vec2 = vec1 @ mats[1]
+    vec3 = vec2 @ mats[2]
+    ref = Tensor.stack(vec1, vec2, vec3)
+    ref.realize()
+    return vec, mats, ref
+
+  def test_uop_scan_matmul(self):
+    vec, mats, ref = self._test_scan()
+
+    # 3 matmuls with SCAN
+    i = UOp.range(3, -100, AxisType.OUTER)
+    out = Tensor.empty(3, 1, 10)
+    phi = Tensor(i.eq(0).where(vec.uop, out[(i-1).maximum(0)].uop))
+    comp = phi @ mats[i]
+    store = out[i].uop.store(comp.uop).end(i)
+    out = Tensor(out.uop.after(store))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+
 class TestOuterworld(unittest.TestCase):
   def test_range_plus_1(self):
     t = Tensor.arange(100).reshape(10,10).realize()
@@ -70,5 +146,85 @@ class TestOuterworld(unittest.TestCase):
     out = out.reshape(1, 3).expand(a, 3).contiguous().realize()
     self.assertListEqual([[0,4,8],[4,8,12],[8,12,16]], out.tolist())
 
+class TestVmap(unittest.TestCase):
+  def test_vmap_inner(self, axis_type=AxisType.LOOP, fuse=False, grad=False):
+    x = Tensor.ones(1, 10).contiguous().requires_grad_()
+    mats = Tensor.ones(3, 10, 10).contiguous().requires_grad_()
+
+    ref = x @ mats
+    if fuse: ref = ref * 2
+
+    # vmap across axis 0
+    a = UOp.range(3, -1, axis_type)
+    out = x @ mats[a]
+    out = out.reshape(1, 10).pad(((a,(3-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    if fuse: out = out * 2
+    if grad:
+      out.mean().backward()
+      np.testing.assert_allclose(mats.grad.numpy(), (2./30) if fuse else (1./30))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+  def test_vmap_inner_fuse(self): self.test_vmap_inner(fuse=True)
+  def test_vmap_outer(self): self.test_vmap_inner(AxisType.OUTER)
+  def test_vmap_outer_fuse(self): self.test_vmap_inner(AxisType.OUTER, fuse=True)
+
+  def test_vmap_inner_grad(self): self.test_vmap_inner(grad=True)
+  def test_vmap_inner_fuse_grad(self): self.test_vmap_inner(fuse=True, grad=True)
+  def test_vmap_outer_grad(self): self.test_vmap_inner(AxisType.OUTER, grad=True)
+
+  def test_vmap_convs(self):
+    layers = [
+      nn.Conv2d(1, 8, 3), Tensor.relu,
+      nn.Conv2d(8, 8, 3), Tensor.relu]
+    img = Tensor.randn(4, 1, 16, 16).realize(*nn.state.get_parameters(layers))
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None, None, None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.realize()
+    np.testing.assert_allclose(out.numpy(), img.sequential(layers).numpy(), atol=1e-6)
+
+  def test_vmap_gemm(self):
+    layers = [
+      nn.Linear(16, 16, bias=False), Tensor.relu,
+      nn.Linear(16, 16, bias=False), Tensor.relu]
+    img = Tensor.randn(4, 16).realize(*nn.state.get_parameters(layers))
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.realize()
+    np.testing.assert_allclose(out.numpy(), img.sequential(layers).numpy(), atol=1e-6)
+
+  @unittest.skip("this is broken, we need to lower the outer reduce in the outer graph")
+  def test_vmap_gemm_grad(self):
+    layers = [
+      nn.Linear(16, 16, bias=False), Tensor.relu,
+      nn.Linear(16, 16, bias=False), Tensor.relu]
+    layer_tensors = nn.state.get_parameters(layers)
+    img = Tensor.randn(4, 16).realize(*layer_tensors)
+    for l in layer_tensors: l.requires_grad_()
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.mean().backward()
+    grads = [l.grad for l in layer_tensors]
+    out.realize(*grads)
+    out_grads = [x.numpy() for x in grads]
+
+    # compute reference grads
+    for l in layer_tensors: l.grad = None
+    img.sequential(layers).mean().backward()
+    grads = [l.grad for l in layer_tensors]
+    out.realize(*grads)
+    ref_grads = [x.numpy() for x in grads]
+
+    # compare
+    for o,r in zip(out_grads, ref_grads): np.testing.assert_allclose(o, r, atol=1e-6)
+
 if __name__ == '__main__':
   unittest.main()

test/test_profiler.py

@@ -17,7 +17,7 @@ def helper_collect_profile(*devs):
   cpu_events.clear()
 
   profile_list = []
-  with Context(VIZ=1):
+  with Context(VIZ=1, PROFILE=1):
     yield profile_list
     for dev in devs: dev.synchronize()
     for dev in devs: dev._at_profile_finalize()

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/test_uops.py

@@ -517,7 +517,7 @@ class TestUOpStr(unittest.TestCase):
 
 class TestUPatHelpers(unittest.TestCase):
   def test_location(self):
-    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "math.py")
+    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "symbolic.py")
     self.assertEqual(shared_spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
     test_upat = UPat(Ops.CONST, dtypes.bool)
     self.assertEqual(test_upat.location[0].split("/")[-1], __file__.replace("\\", "/").split("/")[-1])

test/testextra/test_tk.py

@@ -9,7 +9,7 @@ import numpy as np
 from extra.thunder.tiny.tk import WARP_THREADS
 from extra.thunder.tiny.tk.kernel import Kernel
 
-@unittest.skipUnless(Device.DEFAULT in ["CUDA", "NV"], "only cuda")
+@unittest.skipIf(CI and Device.DEFAULT not in ["CUDA", "NV"], "only cuda")
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "no ptx")
 class TestTK(unittest.TestCase):
   @unittest.skipIf(CI, "no wmma in ci")
@@ -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

@@ -5,16 +5,16 @@ from tinygrad.runtime.support.c import Struct
 class TestAutogen(unittest.TestCase):
   def test_packed_struct_sizeof(self):
     layout = [('a', ctypes.c_char), ('b', ctypes.c_int, 5), ('c', ctypes.c_char)]
-    class X(ctypes.Structure): _fields_, _layout_ = layout, 'gcc-sysv'
     class Y(ctypes.Structure): _fields_, _pack_, _layout_ = layout, 1, 'ms'
-    class Z(Struct): _packed_, _fields_ = True, layout
-    self.assertNotEqual(ctypes.sizeof(X), 4) # ctypes bug! gcc-13.3.0 says this should have size 4
+    class Z(Struct): pass
+    Z._packed_, Z._fields_ = True, layout
     self.assertEqual(ctypes.sizeof(Y), 6)
     self.assertEqual(ctypes.sizeof(Z), 3)
     layout = [('a', ctypes.c_int, 31), ('b', ctypes.c_int, 31), ('c', ctypes.c_int, 1), ('d', ctypes.c_int, 1)]
     class Foo(ctypes.Structure): _fields_, _layout_ = layout, 'gcc-sysv'
     class Bar(ctypes.Structure): _fields_, _pack_, _layout_ = layout, 1, 'ms'
-    class Baz(Struct): _fields_, _packed_ = layout, True
+    class Baz(Struct): pass
+    Baz._packed_, Baz._fields_ = True, layout
     self.assertEqual(ctypes.sizeof(Foo), 12)
     self.assertEqual(ctypes.sizeof(Bar), 12)
     self.assertEqual(ctypes.sizeof(Baz), 8)
@@ -44,4 +44,42 @@ class TestAutogen(unittest.TestCase):
       test.argtypes = [Baz]
       self.assertEqual(test(b), b.a + b.b + b.c + b.d)
 
+  @unittest.skipIf(WIN, "doesn't compile on windows")
+  def test_packed_structs(self):
+    NvU32 = ctypes.c_uint32
+    NvU64 = ctypes.c_uint64
+    class FWSECLIC_READ_VBIOS_DESC(Struct): pass
+    FWSECLIC_READ_VBIOS_DESC._packed_ = True
+    FWSECLIC_READ_VBIOS_DESC._fields_ = [
+      ('version', NvU32),
+      ('size', NvU32),
+      ('gfwImageOffset', NvU64),
+      ('gfwImageSize', NvU32),
+      ('flags', NvU32),
+    ]
+    class FWSECLIC_FRTS_REGION_DESC(Struct): pass
+    FWSECLIC_FRTS_REGION_DESC._packed_ = True
+    FWSECLIC_FRTS_REGION_DESC._fields_ = [
+      ('version', NvU32),
+      ('size', NvU32),
+      ('frtsRegionOffset4K', NvU32),
+      ('frtsRegionSize', NvU32),
+      ('frtsRegionMediaType', NvU32),
+    ]
+    class FWSECLIC_FRTS_CMD(Struct): pass
+    FWSECLIC_FRTS_CMD._packed_ = True
+    FWSECLIC_FRTS_CMD._fields_ = [
+      ('readVbiosDesc', FWSECLIC_READ_VBIOS_DESC),
+      ('frtsRegionDesc', FWSECLIC_FRTS_REGION_DESC),
+    ]
+    read_vbios_desc = FWSECLIC_READ_VBIOS_DESC(version=0x1, size=ctypes.sizeof(FWSECLIC_READ_VBIOS_DESC), flags=2)
+    frst_reg_desc = FWSECLIC_FRTS_REGION_DESC(version=0x1, size=ctypes.sizeof(FWSECLIC_FRTS_REGION_DESC),
+      frtsRegionOffset4K=0xdead, frtsRegionSize=0x100, frtsRegionMediaType=2)
+    frts_cmd = FWSECLIC_FRTS_CMD(readVbiosDesc=read_vbios_desc, frtsRegionDesc=frst_reg_desc)
+    assert int.from_bytes(frts_cmd, 'little') == 0x2000001000000dead0000001400000001000000020000000000000000000000000000001800000001
+    assert int.from_bytes(frts_cmd.readVbiosDesc, 'little') == int.from_bytes(read_vbios_desc, 'little')
+    assert int.from_bytes(frts_cmd.frtsRegionDesc, 'little') == int.from_bytes(frst_reg_desc, 'little')
+    assert frts_cmd.readVbiosDesc.__class__ is FWSECLIC_READ_VBIOS_DESC
+    assert frts_cmd.frtsRegionDesc.__class__ is FWSECLIC_FRTS_REGION_DESC
+
 if __name__ == "__main__": unittest.main()

test/unit/test_microbenchmarks.py

@@ -1,4 +1,5 @@
 import unittest, time
+from tinygrad.helpers import Profiling
 from tinygrad.uop.ops import UOp
 from tinygrad.dtype import dtypes
 
@@ -38,6 +39,14 @@ class TestMicrobenchmarks(unittest.TestCase):
     a = UOp.const(dtypes.int, 2)
     for _ in range(N): (a+a).simplify()
 
+class TestMicroprofile(unittest.TestCase):
+  def test_uop_simplify_complex(self):
+    x = UOp.variable("x", 0, 10)
+    y = UOp.variable("y", 0, 10)
+    expr = (x*2)+5+(x*4)+(y*2)+y
+    with Profiling():
+      for _ in range(1000): expr.simplify()
+
 if __name__ == '__main__':
   unittest.main()
 

test/unit/test_pattern_matcher.py

@@ -211,39 +211,5 @@ class TestPatternMatcher(unittest.TestCase):
       return u.src[0]
     for a,b in zip(simple_src(a), simple_src(b)): self._assert_eq_upat(a, b)
 
-class TestAlgebraic(unittest.TestCase):
-  def test_plus_0(self):
-    pm = PatternMatcher([
-      (UPat.var("x") + 0, UPat.var("x")), # x+0 -> x
-    ])
-    expr = UOp.const(dtypes.int, 4)+0
-    print(expr)
-    self.assertEqual(pm.rewrite(expr), UOp.const(dtypes.int, 4))
-
-  def test_div_mul(self):
-    pm = PatternMatcher([
-      ((UPat.var("x") * UPat.var("x2")) / UPat.var("x2"), UPat.var("x")), # (x*x2)/x2 -> x
-    ])
-    expr = UOp.const(dtypes.float, 4)/2*2
-    print(expr)
-    self.assertEqual(pm.rewrite(expr), UOp.const(dtypes.int, 4))
-
-  def test_mul_is_and(self):
-    pm = PatternMatcher([
-      (UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool), UPat.var('x') & UPat.var('y')),
-    ])
-    expr = UOp.const(dtypes.bool, True)*UOp.const(dtypes.bool, True)
-    print(expr)
-    self.assertEqual(pm.rewrite(expr), UOp.const(dtypes.bool, True)&UOp.const(dtypes.bool, True))
-
-  def test_div_neg_1(self):
-    pm = PatternMatcher([
-      (UPat.var("x") // -1, UPat.var("x") * -1), # x//-1 -> x * -1
-    ])
-    expr = UOp.const(dtypes.float, 4)//-1
-    print(expr)
-    self.assertEqual(pm.rewrite(expr), UOp.const(dtypes.int, 4) * -1)
-
-
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/unit/test_tqdm.py

@@ -66,6 +66,7 @@ class TestProgressBar(unittest.TestCase):
       tqdm_output = tqdm.format_meter(n=total, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
       self._compare_bars(tinytqdm_output, tqdm_output)
 
+  @unittest.skip("this is flaky")
   @patch('sys.stderr', new_callable=StringIO)
   @patch('shutil.get_terminal_size')
   def test_unit_scale(self, mock_terminal_size, mock_stderr):

test/unit/test_viz.py

@@ -6,6 +6,7 @@ from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatch
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
+from tinygrad.helpers import VIZ
 from tinygrad.device import Buffer
 
 @track_rewrites(name=True)
@@ -33,11 +34,14 @@ class BaseTestViz(unittest.TestCase):
     cpu_events.clear()
     self.tms = TRACK_MATCH_STATS.value
     self.profile = PROFILE.value
+    self.viz = VIZ.value
     TRACK_MATCH_STATS.value = 2
     PROFILE.value = 1
+    VIZ.value = 1
   def tearDown(self):
     TRACK_MATCH_STATS.value = self.tms
     PROFILE.value = self.profile
+    VIZ.value = self.viz
 
 class TestViz(BaseTestViz):
   def test_simple(self):

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,15 +13,12 @@ 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
     self.dont_use_locals = self.ast.arg.dont_use_locals if self.ast.arg is not None else False
     self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else []
+    self.opt_range = itertools.count(start=max([x.arg[0] for x in self.rngs], default=0)+1)
 
   @property
   def rngs(self):
@@ -32,8 +30,6 @@ class Scheduler:
   def full_shape(self): return [ssimplify(x.src[0]) for x in self.rngs]
   @property
   def axis_types(self): return [x.arg[-1] for x in self.rngs]
-  @property
-  def maxarg(self): return max([x.arg[0] for x in self.rngs], default=0)
 
   # strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2']
   def shape_str(self) -> list[str]:
@@ -55,8 +51,10 @@ class Scheduler:
   def get_optimized_ast(self, name_override:str|None=None):
     if name_override is not None: name = name_override
     else:
-      kernel_type = "r" if self.reduceop is not None else "E"
-      name = kernel_type + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
+      k_type = "r" if self.reduceop is not None else "E"
+      special_uops = sorted([x for x in self.ast.toposort() if x.op is Ops.SPECIAL], key=lambda x: x.arg)
+      special_ops = [colored(str(x.vmax+1), "blue" if x.arg[0] == "g" else "cyan") for x in special_uops]
+      name = k_type + colored('_', 'BLACK').join(['']+special_ops+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
       Scheduler.kernel_cnt[(function_name := to_function_name(name))] += 1
       num = f"n{Scheduler.kernel_cnt[function_name]-1}" if Scheduler.kernel_cnt[function_name] > 1 else ""
       name += colored(num, 'BLACK')
@@ -96,7 +94,7 @@ class Scheduler:
   def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng=None):
     if (old_sz:=rng.src[0].divides(amount)) is None:
       raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}")
-    new_rng = UOp.range(amount, self.maxarg+1, new_type) if input_new_rng is None else input_new_rng
+    new_rng = UOp.range(amount, next(self.opt_range), new_type) if input_new_rng is None else input_new_rng
     replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
     sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
     self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}")
@@ -232,9 +230,9 @@ class Scheduler:
       for tc in tensor_cores:
         if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
           # tensor cores have three ranges. X, Y, and REDUCE
-          in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: -x.arg[0])
-          in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: -x.arg[0])
-          red_ranges = sorted(reduceop.src[1:], key=lambda x: -x.arg[0])
+          in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0], reverse=True)
+          in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0], reverse=True)
+          red_ranges = sorted(reduceop.src[1:], key=lambda x: x.arg[0], reverse=True)
           if DEBUG >= 3:
             print(f"TC({axis}): {[(x.arg[0],x.vmax+1) for x in in0_ranges]}",
                               f"{[(x.arg[0],x.vmax+1) for x in in1_ranges]} {[(x.arg[0],x.vmax+1) for x in red_ranges]}")

tinygrad/device.py

@@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
-from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM, select_first_inited
+from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM, select_first_inited, VIZ
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer
 
@@ -355,8 +355,9 @@ if PROFILE:
 
     with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(cpu_events+Compiled.profile_events+Buffer.profile_events, f)
 
-    from tinygrad.uop.ops import launch_viz
-    launch_viz("PROFILE", fn)
+    if VIZ:
+      from tinygrad.uop.ops import launch_viz
+      launch_viz("PROFILE", fn)
 
 def enumerate_devices_str() -> Generator[str, None, None]:
   from tinygrad import Tensor, Device

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
 
@@ -22,10 +23,13 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
   in_degree: dict[UOp, int] = {}
   var_vals: dict[str, int] = {}
   for u in sched_sink.toposort():
-    if u.op is not Ops.AFTER: continue  # anything that's not an ASSIGN doesn't write a kernel, so we can skip
+    if u.op is Ops.RANGE:
+      in_degree.setdefault(u, 0)
+      continue
+    if u.op is not Ops.AFTER or u.src[1].op is Ops.RANGE: continue
     k = u.src[1]
     in_degree.setdefault(k, 0)
-    for s in k.src:
+    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 +43,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 +64,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))
-    for x in children[k]:
+      raise RuntimeError(f"can't schedule {k.op}")
+    for x in children[rk]:
       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

@@ -3,14 +3,15 @@ import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
 from tinygrad.helpers import argsort
 
-def reduce_gradient(ctx:UOp, ret:UOp):
+def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   def broadcast_to_input(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)-len(x.shape))).expand(ret.src[0].shape)
-  if ret.arg[0] == Ops.ADD: return (broadcast_to_input(ctx),)
-  if ret.arg[0] == Ops.MAX:
+  if op == Ops.ADD: return (broadcast_to_input(ctx),)
+  if op == Ops.MAX:
+    assert ret.op is Ops.REDUCE_AXIS, "only works on REDUCE_AXIS"
     mask = ret.src[0].eq(broadcast_to_input(ret)).cast(ctx.dtype)
     count = mask.r(Ops.ADD, ret.arg[1])
     return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
-  if ret.arg[0] == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
+  if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
 
 # ctx is grad_output
 pm_gradient = PatternMatcher([
@@ -28,7 +29,8 @@ pm_gradient = PatternMatcher([
     ((x>y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)), (x<y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)))),
   (UPat(Ops.MUL, name="ret"), lambda ctx, ret: (ret.src[1]*ctx, ret.src[0]*ctx)),
   (UPat(Ops.WHERE, name="ret"), lambda ctx, ret: (None, ret.src[0].where(ctx, ctx.const_like(0)), ret.src[0].where(ctx.const_like(0), ctx))),
-  (UPat(Ops.REDUCE_AXIS, name="ret"), reduce_gradient),
+  (UPat(Ops.REDUCE_AXIS, name="ret"), lambda ctx, ret: reduce_gradient(ctx, ret, ret.arg[0])),
+  (UPat(Ops.REDUCE, name="ret"), lambda ctx, ret: reduce_gradient(ctx, ret, ret.arg) + (None,)*(len(ret.src)-1)),
   (UPat(Ops.CONTIGUOUS), lambda ctx: (ctx,)),
   (UPat(Ops.CONTIGUOUS_BACKWARD), lambda ctx: (ctx.contiguous(),)),
   (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)),
@@ -68,4 +70,8 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
         # we add the backward metadata to everything new in the graph
         for bw_uop in v.toposort(lambda x: x not in (t0, *t0.src, grads[t0])):
           all_metadata[bw_uop] = all_metadata.get(bw_uop, ())+backward_metadata
+  # end any ranges on grads with a reduce sum
+  for k,v in grads.items():
+    if len(v.ranges):
+      grads[k] = v.reduce(*v.ranges, arg=Ops.ADD)
   return grads

tinygrad/helpers.py

@@ -179,7 +179,9 @@ ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), Conte
 EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
 CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
-VIZ = PROFILE = ContextVar("VIZ", 0)
+# VIZ implies PROFILE, but you can run PROFILE without VIZ
+VIZ = ContextVar("VIZ", 0)
+PROFILE = ContextVar("PROFILE", VIZ.value)
 SPEC = ContextVar("SPEC", 1)
 # TODO: disable by default due to speed
 IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
@@ -416,6 +418,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,11 +1,11 @@
 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
-import base64, ctypes, ctypes.util, struct, functools, inspect
+from tinygrad.runtime.autogen import mesa
+import base64, contextlib, ctypes, ctypes.util, struct, functools, inspect
 
 def g(s:str): return getattr(mesa, s)
 def nsrc(d:mesa.nir_def) -> mesa.nir_src: return mesa.nir_src(ssa=ctypes.pointer(d))
@@ -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)
@@ -157,8 +157,7 @@ class NIRRenderer(Renderer):
   def __init__(self): mesa.glsl_type_singleton_init_or_ref()
 
   def __del__(self):
-    try: mesa.glsl_type_singleton_decref()
-    except FileNotFoundError: pass
+    with contextlib.suppress(AttributeError):mesa.glsl_type_singleton_decref()
 
   @property
   def nir_options(self): raise NotImplementedError("needs nir_options")
@@ -177,7 +176,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)))