From 03a6bc59c16c609f71bb14b9dcc243cc5f5c4c7b Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Fri, 26 Jan 2024 12:44:19 -0800
Subject: [PATCH] move autogen to runtime/autogen (#3254)

---
 .github/workflows/test.yml               | 16 +++----
 autogen_stubs.sh                         | 54 ++++++++++++------------
 extra/dist/world.py                      |  2 +-
 extra/hip_events.py                      |  2 +-
 ruff.toml                                |  2 +-
 setup.py                                 |  2 +-
 sz.py                                    |  2 +-
 tinygrad/{ => runtime}/autogen/comgr.py  |  0
 tinygrad/{ => runtime}/autogen/cuda.py   |  0
 tinygrad/{ => runtime}/autogen/hip.py    |  0
 tinygrad/{ => runtime}/autogen/opencl.py |  0
 tinygrad/runtime/graph/cuda.py           |  2 +-
 tinygrad/runtime/graph/hip.py            |  2 +-
 tinygrad/runtime/ops_cuda.py             |  2 +-
 tinygrad/runtime/ops_gpu.py              |  2 +-
 tinygrad/runtime/ops_hip.py              |  2 +-
 16 files changed, 46 insertions(+), 44 deletions(-)
 rename tinygrad/{ => runtime}/autogen/comgr.py (100%)
 rename tinygrad/{ => runtime}/autogen/cuda.py (100%)
 rename tinygrad/{ => runtime}/autogen/hip.py (100%)
 rename tinygrad/{ => runtime}/autogen/opencl.py (100%)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 22f1214633..a7bfb78a2f 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -330,11 +330,11 @@ jobs:
       run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
     - name: Verify HIP autogen
       run: |
-        cp tinygrad/autogen/hip.py /tmp/hip.py.bak
-        cp tinygrad/autogen/comgr.py /tmp/comgr.py.bak
+        cp tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
+        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
         ./autogen_stubs.sh hip
-        diff /tmp/hip.py.bak tinygrad/autogen/hip.py
-        diff /tmp/comgr.py.bak tinygrad/autogen/comgr.py
+        diff /tmp/hip.py.bak tinygrad/runtime/autogen/hip.py
+        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
     - name: Test HIP compilation on RDNA3 [gfx1100]
       # test/test_symbolic_ops.py can't run here, it was comparing empty memory
       run: |
@@ -453,15 +453,15 @@ jobs:
       - name: Verify OpenCL autogen
         if: matrix.backend == 'gpu'
         run: |
-          cp tinygrad/autogen/opencl.py /tmp/opencl.py.bak
+          cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
           ./autogen_stubs.sh opencl
-          diff /tmp/opencl.py.bak tinygrad/autogen/opencl.py
+          diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
       - name: Verify CUDA autogen
         if: matrix.backend == 'cuda'
         run: |
-          cp tinygrad/autogen/cuda.py /tmp/cuda.py.bak
+          cp tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
           ./autogen_stubs.sh cuda
-          diff /tmp/cuda.py.bak tinygrad/autogen/cuda.py
+          diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
       - name: Run pytest (not cuda)
         if: matrix.backend!='cuda' && matrix.backend!='ptx' && matrix.backend!='triton'
         run: python -m pytest -n=auto test/ --durations=20
diff --git a/autogen_stubs.sh b/autogen_stubs.sh
index 4ef21f7654..922ac28e7e 100755
--- a/autogen_stubs.sh
+++ b/autogen_stubs.sh
@@ -13,47 +13,49 @@ if [[ ! $(clang2py -V) ]]; then
   popd
 fi
 
+BASE=tinygrad/runtime/autogen/
+
 generate_opencl() {
-  clang2py /usr/include/CL/cl.h -o autogen/opencl.py -l /usr/lib/x86_64-linux-gnu/libOpenCL.so.1 -k cdefstum
-  sed -i '1s/^/# mypy: ignore-errors\n/' autogen/opencl.py
-  grep FIXME_STUB autogen/opencl.py || true
+  clang2py /usr/include/CL/cl.h -o $BASE/opencl.py -l /usr/lib/x86_64-linux-gnu/libOpenCL.so.1 -k cdefstum
+  sed -i '1s/^/# mypy: ignore-errors\n/' $BASE/opencl.py
+  grep FIXME_STUB $BASE/opencl.py || true
   # hot patches
-  sed -i "s\import ctypes\import ctypes, ctypes.util\g" autogen/opencl.py
-  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libOpenCL.so.1')\ctypes.CDLL(ctypes.util.find_library('OpenCL'))\g" autogen/opencl.py
-  python3 -c "import autogen.opencl"
+  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 autogen/hip.py -l /opt/rocm/lib/libamdhip64.so
-  echo "hipDeviceProp_t = hipDeviceProp_tR0600" >> autogen/hip.py
-  echo "hipGetDeviceProperties = hipGetDevicePropertiesR0600" >> autogen/hip.py
-  sed -i '1s/^/# mypy: ignore-errors\n/' autogen/hip.py
-  grep FIXME_STUB autogen/hip.py || true
+  --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
+  sed -i '1s/^/# mypy: ignore-errors\n/' $BASE/hip.py
+  grep FIXME_STUB $BASE/hip.py || true
   # we can trust HIP is always at /opt/rocm/lib
-  #sed -i "s\import ctypes\import ctypes, ctypes.util\g" autogen/hip.py
-  #sed -i "s\ctypes.CDLL('/opt/rocm/lib/libhiprtc.so')\ctypes.CDLL(ctypes.util.find_library('hiprtc'))\g" autogen/hip.py
-  #sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamdhip64.so')\ctypes.CDLL(ctypes.util.find_library('amdhip64'))\g" autogen/hip.py
-  python3 -c "import autogen.hip"
+  #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
+  python3 -c "import tinygrad.runtime.autogen.hip"
 
   clang2py /opt/rocm/include/amd_comgr/amd_comgr.h \
-  --clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o autogen/comgr.py -l /opt/rocm/lib/libamd_comgr.so
-  sed -i '1s/^/# mypy: ignore-errors\n/' autogen/comgr.py
-  grep FIXME_STUB autogen/comgr.py || true
-  python3 -c "import autogen.comgr"
+  --clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/comgr.py -l /opt/rocm/lib/libamd_comgr.so
+  sed -i '1s/^/# mypy: ignore-errors\n/' $BASE/comgr.py
+  grep FIXME_STUB $BASE/comgr.py || true
+  python3 -c "import tinygrad.runtime.autogen.comgr"
 }
 
 generate_cuda() {
-  clang2py /usr/include/cuda.h /usr/include/nvrtc.h -o autogen/cuda.py -l /usr/lib/x86_64-linux-gnu/libcuda.so -l /usr/lib/x86_64-linux-gnu/libnvrtc.so
-  sed -i "s\import ctypes\import ctypes, ctypes.util\g" autogen/cuda.py
-  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libcuda.so')\ctypes.CDLL(ctypes.util.find_library('cuda'))\g" autogen/cuda.py
-  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libnvrtc.so')\ctypes.CDLL(ctypes.util.find_library('nvrtc'))\g" autogen/cuda.py
-  sed -i '1s/^/# mypy: ignore-errors\n/' autogen/cuda.py
-  grep FIXME_STUB autogen/cuda.py || true
+  clang2py /usr/include/cuda.h /usr/include/nvrtc.h -o $BASE/cuda.py -l /usr/lib/x86_64-linux-gnu/libcuda.so -l /usr/lib/x86_64-linux-gnu/libnvrtc.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
+  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libnvrtc.so')\ctypes.CDLL(ctypes.util.find_library('nvrtc'))\g" $BASE/cuda.py
+  sed -i '1s/^/# mypy: ignore-errors\n/' $BASE/cuda.py
+  grep FIXME_STUB $BASE/cuda.py || true
+  python3 -c "import tinygrad.runtime.autogen.cuda"
 }
 
-cd tinygrad
 if [ "$1" == "opencl" ]; then generate_opencl
 elif [ "$1" == "hip" ]; then generate_hip
 elif [ "$1" == "cuda" ]; then generate_cuda
diff --git a/extra/dist/world.py b/extra/dist/world.py
index 997082cd98..79ee0d7cdf 100644
--- a/extra/dist/world.py
+++ b/extra/dist/world.py
@@ -5,7 +5,7 @@ from tinygrad.helpers import DEBUG, colored, getenv
 from tinygrad.lazy import LazyBuffer
 from tinygrad.runtime.lib import RawBuffer, RawBufferCopyInOut
 try:
-  import tinygrad.autogen.hip as hip
+  import tinygrad.runtime.autogen.hip as hip
   from tinygrad.runtime.ops_hip import RawHIPBuffer, check
 except: RawHIPBuffer = None
 from tinygrad.runtime.ops_disk import RawDiskBuffer
diff --git a/extra/hip_events.py b/extra/hip_events.py
index 7a5e432f15..5719a18ce9 100644
--- a/extra/hip_events.py
+++ b/extra/hip_events.py
@@ -1,5 +1,5 @@
 import ctypes
-import tinygrad.autogen.hip as hip
+import tinygrad.runtime.autogen.hip as hip
 from tinygrad.runtime.ops_hip import check
 from tinygrad.helpers import init_c_var
 
diff --git a/ruff.toml b/ruff.toml
index 6b91b45852..bbe3319a98 100644
--- a/ruff.toml
+++ b/ruff.toml
@@ -31,6 +31,6 @@ exclude = [
   "examples/",
   "extra/",
   "openpilot/",
-  "tinygrad/autogen",
+  "tinygrad/runtime/autogen",
 ]
 
diff --git a/setup.py b/setup.py
index 765ec77a8d..bdab81dab6 100644
--- a/setup.py
+++ b/setup.py
@@ -14,7 +14,7 @@ setup(name='tinygrad',
       license='MIT',
       long_description=long_description,
       long_description_content_type='text/markdown',
-      packages = ['tinygrad', 'tinygrad.autogen', 'tinygrad.codegen', 'tinygrad.nn', 'tinygrad.renderer',
+      packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.codegen', 'tinygrad.nn', 'tinygrad.renderer',
                   'tinygrad.runtime', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.features'],
       classifiers=[
         "Programming Language :: Python :: 3",
diff --git a/sz.py b/sz.py
index 67a409743b..aaec10b21c 100755
--- a/sz.py
+++ b/sz.py
@@ -12,7 +12,7 @@ def gen_stats(base_path="."):
   for path, _, files in os.walk(os.path.join(base_path, "tinygrad")):
     for name in files:
       if not name.endswith(".py"): continue
-      if 'tinygrad/autogen' in path: continue
+      if 'tinygrad/runtime/autogen' in path: continue
       filepath = os.path.join(path, name)
       relfilepath = os.path.relpath(filepath, base_path)
       with tokenize.open(filepath) as file_:
diff --git a/tinygrad/autogen/comgr.py b/tinygrad/runtime/autogen/comgr.py
similarity index 100%
rename from tinygrad/autogen/comgr.py
rename to tinygrad/runtime/autogen/comgr.py
diff --git a/tinygrad/autogen/cuda.py b/tinygrad/runtime/autogen/cuda.py
similarity index 100%
rename from tinygrad/autogen/cuda.py
rename to tinygrad/runtime/autogen/cuda.py
diff --git a/tinygrad/autogen/hip.py b/tinygrad/runtime/autogen/hip.py
similarity index 100%
rename from tinygrad/autogen/hip.py
rename to tinygrad/runtime/autogen/hip.py
diff --git a/tinygrad/autogen/opencl.py b/tinygrad/runtime/autogen/opencl.py
similarity index 100%
rename from tinygrad/autogen/opencl.py
rename to tinygrad/runtime/autogen/opencl.py
diff --git a/tinygrad/runtime/graph/cuda.py b/tinygrad/runtime/graph/cuda.py
index e345c90117..9eca6eaa05 100644
--- a/tinygrad/runtime/graph/cuda.py
+++ b/tinygrad/runtime/graph/cuda.py
@@ -1,6 +1,6 @@
 import ctypes
 from typing import Any, Optional, Tuple, Dict, List, cast
-import tinygrad.autogen.cuda as cuda
+import tinygrad.runtime.autogen.cuda as cuda
 from tinygrad.helpers import init_c_var, encode_args_cuda_style, all_same, GraphException
 from tinygrad.device import CompiledASTRunner, update_stats, Buffer
 from tinygrad.runtime.ops_cuda import check, cu_time_execution
diff --git a/tinygrad/runtime/graph/hip.py b/tinygrad/runtime/graph/hip.py
index 237765aa8f..54e186e821 100644
--- a/tinygrad/runtime/graph/hip.py
+++ b/tinygrad/runtime/graph/hip.py
@@ -1,6 +1,6 @@
 import ctypes
 from typing import Tuple
-import tinygrad.autogen.hip as hip
+import tinygrad.runtime.autogen.hip as hip
 from tinygrad.helpers import init_c_var, time_execution_cuda_style
 from tinygrad.runtime.ops_hip import check
 from tinygrad.runtime.graph.cuda import CUDAGraph
diff --git a/tinygrad/runtime/ops_cuda.py b/tinygrad/runtime/ops_cuda.py
index 1e9beb23e6..f4451508b1 100644
--- a/tinygrad/runtime/ops_cuda.py
+++ b/tinygrad/runtime/ops_cuda.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 import subprocess, hashlib, tempfile, ctypes, ctypes.util, functools, re
 from pathlib import Path
 from typing import Tuple, Optional
-import tinygrad.autogen.cuda as cuda
+import tinygrad.runtime.autogen.cuda as cuda
 from tinygrad.helpers import DEBUG, getenv, from_mv, init_c_var, colored, cpu_time_execution, compile_cuda_style, encode_args_cuda_style, time_execution_cuda_style  # noqa: E501
 from tinygrad.device import Compiled, LRUAllocator, MallocAllocator
 from tinygrad.codegen.kernel import LinearizerOptions
diff --git a/tinygrad/runtime/ops_gpu.py b/tinygrad/runtime/ops_gpu.py
index dff60a2af2..1da6b397ed 100644
--- a/tinygrad/runtime/ops_gpu.py
+++ b/tinygrad/runtime/ops_gpu.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 from typing import Tuple, Optional, List, cast
 import ctypes, functools, hashlib
-import tinygrad.autogen.opencl as cl
+import tinygrad.runtime.autogen.opencl as cl
 from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG
 from tinygrad.codegen.kernel import LinearizerOptions
 from tinygrad.renderer.cstyle import OpenCLRenderer
diff --git a/tinygrad/runtime/ops_hip.py b/tinygrad/runtime/ops_hip.py
index 85d0f7b040..d311cfad4e 100644
--- a/tinygrad/runtime/ops_hip.py
+++ b/tinygrad/runtime/ops_hip.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 import ctypes, functools, subprocess, io
 from typing import Tuple, TypeVar, List, Any, cast, Set
-import tinygrad.autogen.hip as hip
+import tinygrad.runtime.autogen.hip as hip
 from tinygrad.helpers import DEBUG, getenv, init_c_var
 from tinygrad.helpers import from_mv, round_up, to_mv, colored, init_c_struct_t, to_char_p_p, get_bytes
 from tinygrad.device import Compiled, LRUAllocator, MallocAllocator, BufferOptions, JITRunner, Device, Buffer, update_stats