move graph to runtime, check line count with sz.py (#2842)

* move graph to runtime, check line count with sz.py

* oops, didn't save

* dtype aliases

* restore comment, REALCOUNT

.github/workflows/test.yml

@@ -31,8 +31,8 @@ jobs:
         key: linting-packages-${{ hashFiles('**/setup.py') }}-3.8
     - name: Install dependencies
       run: pip install -e '.[linting,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
-    - name: Repo line count
-      run: python sz.py
+    - name: Repo line count <5000 lines
+      run: MAX_LINE_COUNT=5000 python sz.py
     - name: Lint with pylint
       run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' **/*.py
     - name: Lint with ruff
@@ -43,10 +43,6 @@ jobs:
       run: python -m pylint tinygrad/
     - name: Run mypy
       run: python -m mypy
-    - name: Install SLOCCount
-      run: sudo apt install sloccount
-    - name: Check <5000 lines
-      run: sloccount tinygrad test examples extra; if [ $(sloccount tinygrad | sed -n 's/.*Total Physical Source Lines of Code (SLOC)[ ]*= \([^ ]*\).*/\1/p' | tr -d ',') -gt 5000 ]; then exit 1; fi
     - name: Test Docs
       run: |
         python docs/abstractions.py

setup.py

@@ -15,7 +15,7 @@ setup(name='tinygrad',
       long_description=long_description,
       long_description_content_type='text/markdown',
       packages = ['tinygrad', 'tinygrad.codegen', 'tinygrad.nn', 'tinygrad.renderer',
-                  'tinygrad.runtime', 'tinygrad.shape', 'tinygrad.features', 'tinygrad.features.graph'],
+                  'tinygrad.runtime', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.features'],
       classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: MIT License"

sz.py

@@ -4,6 +4,7 @@ import token
 import tokenize
 import itertools
 from tabulate import tabulate
+from tinygrad.helpers import getenv
 
 TOKEN_WHITELIST = [token.OP, token.NAME, token.NUMBER, token.STRING]
 
@@ -16,7 +17,10 @@ def gen_stats(base_path="."):
       relfilepath = os.path.relpath(filepath, base_path)
       with tokenize.open(filepath) as file_:
         tokens = [t for t in tokenize.generate_tokens(file_.readline) if t.type in TOKEN_WHITELIST]
-        token_count, line_count = len(tokens), len(set([t.start[0] for t in tokens]))
+        if getenv("REALCOUNT"):
+          token_count, line_count = len(tokens), len(set([x for t in tokens for x in range(t.start[0], t.end[0]+1)]))
+        else:
+          token_count, line_count = len(tokens), len(set([t.start[0] for t in tokens]))
         table.append([relfilepath, line_count, token_count/line_count])
   return table
 
@@ -65,4 +69,7 @@ if __name__ == "__main__":
       print(tabulate([headers] + sorted(table, key=lambda x: -x[1]), headers="firstrow", floatfmt=".1f")+"\n")
       for dir_name, group in itertools.groupby(sorted([(x[0].rsplit("/", 1)[0], x[1], x[2]) for x in table]), key=lambda x:x[0]):
         print(f"{dir_name:30s} : {sum([x[1] for x in group]):6d}")
-      print(f"\ntotal line count: {sum([x[1] for x in table])}")
+      total_lines = sum([x[1] for x in table])
+      print(f"\ntotal line count: {total_lines}")
+      max_line_count = getenv("MAX_LINE_COUNT", -1)
+      assert max_line_count == -1 or total_lines < max_line_count, f"OVER {max_line_count} LINES"

tinygrad/helpers.py

@@ -150,32 +150,24 @@ class dtypes:
   @staticmethod
   def fields() -> Dict[str, DType]: return DTYPES_DICT
   bool: Final[DType] = DType(0, 1, "bool", np.bool_)
-  float16: Final[DType] = DType(9, 2, "half", np.float16)
-  half = float16
-  float32: Final[DType] = DType(11, 4, "float", np.float32)
-  float = float32
-  float64: Final[DType] = DType(12, 8, "double", np.float64)
-  double = float64
   int8: Final[DType] = DType(1, 1, "char", np.int8)
-  char = int8
-  int16: Final[DType] = DType(3, 2, "short", np.int16)
-  short = int16
-  int32: Final[DType] = DType(5, 4, "int", np.int32)
-  int = int32
-  int64: Final[DType] = DType(7, 8, "long", np.int64)
-  long = int64
   uint8: Final[DType] = DType(2, 1, "unsigned char", np.uint8)
-  uchar = uint8
+  int16: Final[DType] = DType(3, 2, "short", np.int16)
   uint16: Final[DType] = DType(4, 2, "unsigned short", np.uint16)
-  ushort = uint16
+  int32: Final[DType] = DType(5, 4, "int", np.int32)
   uint32: Final[DType] = DType(6, 4, "unsigned int", np.uint32)
-  uint = uint32
+  int64: Final[DType] = DType(7, 8, "long", np.int64)
   uint64: Final[DType] = DType(8, 8, "unsigned long", np.uint64)
-  ulong = uint64
-
-  # NOTE: bfloat16 isn't supported in numpy
-  # it has higher priority than float16, so least_upper_dtype(dtypes.int64, dtypes.uint64) = dtypes.float16
+  float16: Final[DType] = DType(9, 2, "half", np.float16)
+  # bfloat16 has higher priority than float16, so least_upper_dtype(dtypes.int64, dtypes.uint64) = dtypes.float16
   bfloat16: Final[DType] = DType(10, 2, "__bf16", None)
+  float32: Final[DType] = DType(11, 4, "float", np.float32)
+  float64: Final[DType] = DType(12, 8, "double", np.float64)
+
+  # dtype aliases
+  half = float16; float = float32; double = float64 # noqa: E702
+  uchar = uint8; ushort = uint16; uint = uint32; ulong = uint64 # noqa: E702
+  char = int8; short = int16; int = int32; long = int64 # noqa: E702
 
   # NOTE: these are image dtypes
   @staticmethod

tinygrad/runtime/graph/hip.py

@@ -3,7 +3,7 @@ from typing import Tuple
 import gpuctypes.hip as hip
 from tinygrad.helpers import init_c_var
 from tinygrad.runtime.ops_hip import check, hip_time_execution
-from tinygrad.features.graph.cuda import CUDAGraph
+from tinygrad.runtime.graph.cuda import CUDAGraph
 
 class HIPGraph(CUDAGraph):
   def __del__(self):

tinygrad/runtime/ops_cuda.py

@@ -74,7 +74,7 @@ class CUDADevice(Compiled):
       check(cuda.cuDeviceComputeCapability(ctypes.byref(major := ctypes.c_int()), ctypes.byref(minor := ctypes.c_int()), device_id))
       if device_id == 0: CUDADevice.default_arch_name = f"sm_{major.value}{minor.value}"
 
-    from tinygrad.features.graph.cuda import CUDAGraph
+    from tinygrad.runtime.graph.cuda import CUDAGraph
     super().__init__(CUDAAllocator(self) if not CUDACPU else MallocAllocator,
                      LinearizerOptions(supports_float4_alu=False, global_max=[65535, 65535, 2147483647], local_max=[64, 1024, 1024]),
                      CUDARenderer, compile_cuda, functools.partial(CUDAProgram, self), graph=CUDAGraph if not CUDACPU else None)

tinygrad/runtime/ops_hip.py

@@ -66,7 +66,7 @@ class HIPDevice(Compiled):
     self.device = int(device.split(":")[1]) if ":" in device else 0
     if self.device == 0 and not MOCKHIP: HIPDevice.default_arch_name = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device))).gcnArchName.decode()  # noqa: E501
 
-    from tinygrad.features.graph.hip import HIPGraph
+    from tinygrad.runtime.graph.hip import HIPGraph
     super().__init__(MallocAllocator if MOCKHIP else HIPAllocator(self.device), LinearizerOptions(device="HIP"), HIPRenderer,
                      compile_hip, functools.partial(HIPProgram, self.device), HIPGraph)
   def synchronize(self):

tinygrad/runtime/ops_metal.py

@@ -80,7 +80,7 @@ class MetalDevice(Compiled):
     self.mtl_queue = self.device.newCommandQueueWithMaxCommandBufferCount_(1024)
     self.mtl_buffers_in_flight: List[Any] = []
     self.mv_in_metal: List[memoryview] = []
-    from tinygrad.features.graph.metal import MetalGraph
+    from tinygrad.runtime.graph.metal import MetalGraph
     super().__init__(MetalAllocator(self), LinearizerOptions(device="METAL"), MetalRenderer,
                      compile_metal, functools.partial(MetalProgram, self), functools.partial(MetalGraph, self))
   def synchronize(self):