From eb595588bb4cfbf6048b9a2bfa649fba15c4ba2b Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Fri, 1 Dec 2023 14:59:33 -0800
Subject: [PATCH] device.py cleanups for -1 line

---
 tinygrad/device.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/tinygrad/device.py b/tinygrad/device.py
index b68f1a70d4..0a06374e39 100644
--- a/tinygrad/device.py
+++ b/tinygrad/device.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 import numpy as np
 from collections import defaultdict
 from typing import TYPE_CHECKING, Union, Any, List, Optional, Dict, Callable, Tuple
-import importlib, inspect, functools, pathlib, time, re
+import importlib, inspect, functools, pathlib, time, re, ctypes
 from tinygrad.helpers import ansilen, DEBUG, getenv, GlobalCounters, colored, BEAM, NOOPT, all_int, to_function_name, DType, from_mv, dtypes, flat_mv
 from tinygrad.shape.symbolic import Variable, sym_infer, sint
 from tinygrad.ops import LazyOp, TernaryOps, get_lazyop_info, ReduceOps, BufferOps, BinaryOps, Op
@@ -33,6 +33,8 @@ class _Device:
     return "CPU"
 Device = _Device()
 
+# **************** Buffer / Allocator ****************
+
 class Buffer:
   def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None):
     assert isinstance(dtype, DType)
@@ -103,7 +105,14 @@ class LRUAllocator(Allocator):  # pylint: disable=abstract-method
       opaques.clear()
   def free(self, opaque:Any, size:int, dtype:DType): self.cache[(size, dtype)].append(opaque)
 
-# **************** shared device helpers ****************
+class _MallocAllocator(LRUAllocator):
+  def _alloc(self, size:int, dtype:DType): return (ctypes.c_uint8 * (size*dtype.itemsize))()
+  def as_buffer(self, src) -> memoryview: return flat_mv(memoryview(src))
+  def copyin(self, dest, src:memoryview): ctypes.memmove(dest, from_mv(src), len(src))
+  def copyout(self, dest:memoryview, src): ctypes.memmove(from_mv(dest), src, len(dest))
+MallocAllocator = _MallocAllocator()
+
+# **************** base Runner + helpers ****************
 
 class JITRunner:
   def __init__(self):
@@ -271,11 +280,3 @@ def _get_optimized_linearizer(linearizer_opts:LinearizerOptions, ast:LazyOp) ->
       if DEBUG >= 1: print("  <  ".join(f"{nm:6s} : {lin.colored_shape(30, dense=True)} : {tm*1e6:8.2f} us" for nm, lin, tm in timed))
       k = timed[0][1]
   return k
-
-import ctypes
-class _MallocAllocator(LRUAllocator):
-  def _alloc(self, size:int, dtype:DType): return (ctypes.c_uint8 * (size*dtype.itemsize))()
-  def as_buffer(self, src) -> memoryview: return flat_mv(memoryview(src))
-  def copyin(self, dest, src:memoryview): ctypes.memmove(dest, from_mv(src), len(src))
-  def copyout(self, dest:memoryview, src): ctypes.memmove(from_mv(dest), src, len(dest))
-MallocAllocator = _MallocAllocator()