be lazy with the gpubuffer copies for host for constant folding
MODIFY
openpilot/compile.py
openpilot/compile.py
@@ -6,6 +6,7 @@ import os
import time
import io
os.environ['LAZY'] = '1'
+
if int(os.getenv("NOIMAGE", 0)):
pass
else:

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -50,25 +50,26 @@ code_for_op = {
}
class GPUBuffer:
- def __init__(self, shape, hostbuf:Optional[GPUBuffer]=None):
+ def __init__(self, shape, hostbuf:Optional[GPUBuffer]=None, backing:Optional[np.ndarray]=None):
self.st = ShapeTracker(shape)
self.shape = self.st.shape
self._buf : cl.Buffer = hostbuf._buf if hostbuf is not None else None
+ self._base_shape = hostbuf._base_shape if hostbuf is not None else self.shape
+ self._backing = hostbuf._backing if hostbuf is not None else backing
@property
def cl(self):
- if self._buf is None: self._buf = CL.malloc(4*prod(self.shape))
+ if self._buf is None: self._buf = CL.malloc(4*prod(self._base_shape))
+ if self._backing is not None:
+ CL.enqueue_copy(self._buf, self._backing, is_blocking=False)
+ self._backing = None
return self._buf
def __repr__(self): return f"<GPUBuffer with shape {self.shape!r}>"
def shapeTrackerView(self, st:ShapeTracker): return GPUBuffer(st, hostbuf=self)
@staticmethod
- def fromCPU(x):
- ret = GPUBuffer(x.shape)
- # TODO: this is blocking even though we told it not to
- CL.enqueue_copy(ret.cl, x.view(np.ndarray).astype(np.float32).ravel(), is_blocking=False)
- return ret
+ def fromCPU(x): return GPUBuffer(x.shape, backing=x.view(np.ndarray).astype(np.float32).ravel())
def toCPU(self):
data = np.empty(self.shape, dtype=np.float32)

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -2,10 +2,13 @@ from __future__ import annotations
from enum import Enum
from typing import Optional, Tuple, NamedTuple, Union, Any, List, Dict, Type
from copy import copy
-import sys, functools, operator
+import os, sys, functools, operator
from tinygrad.helpers import ConvArgs
from tinygrad.shapetracker import ShapeTracker
+# lazy can recurse a lot
+sys.setrecursionlimit(10000)
+
# these are the llops your accelerator must implement
UnaryOps = Enum("UnaryOps", ["NOOP", "NEG", "RELU", "EXP", "LOG", "SIGN"])
BinaryOps = Enum("BinaryOps", ["ADD", "SUB", "MUL", "DIV", "POW", "CMPEQ"])
@@ -20,24 +23,13 @@ OpType = Union[Type[UnaryOps], Type[BinaryOps], Type[ReduceOps], Type[MovementOp
# TODO: get device buffer types
DeviceBuffer = Any
-# -O1
-MERGE_MOVEMENT_OPS = True
-REMOVE_MOVEMENT_NOPS = True
-MERGE_UNARY_OPS = True
-
-# -O2
-MERGE_ELEMENTWISE_OPS = False # this is making training very slow in the tests
-SHUFFLE_MOVEMENT_OPS = False # this is making training very slow in the tests
-
-# -O3
-SHUFFLE_SLICE_OPS = False # NOTE: 0/0 is NaN if you slice, so this can change the output
-
-# lazy can recurse a lot
-sys.setrecursionlimit(10000)
-
-import os
DEBUG = int(os.getenv("DEBUG", "0"))
GRAPH = int(os.getenv("GRAPH", "0"))
+OPT = int(os.getenv("OPT", "1"))
+
+MERGE_MOVEMENT_OPS, REMOVE_MOVEMENT_NOPS, MERGE_UNARY_OPS = OPT>=1, OPT>=1, OPT>=1
+MERGE_ELEMENTWISE_OPS, SHUFFLE_MOVEMENT_OPS = OPT>=2, OPT>=2
+SHUFFLE_SLICE_OPS = OPT>=3 # NOTE: 0/0 is NaN if you slice, so this can change the output
from collections import defaultdict
cnts : Dict[OpType, int] = defaultdict(int)

be lazy with the gpubuffer copies for host for constant folding