LAZY and CLCACHE are defaults
MODIFY
.github/workflows/test.yml
.github/workflows/test.yml
@@ -51,7 +51,7 @@ jobs:
- name: Install Dependencies
run: pip install -e '.[testing]'
- name: Run Pytest
- run: python -m pytest -s -v
+ run: LAZY=0 python -m pytest -s -v
- name: Run Pytest (lazy)
run: LAZY=1 python -m pytest -s -v
@@ -71,7 +71,7 @@ jobs:
- name: Install Dependencies
run: pip install -e '.[testing]'
- name: Run Pytest
- run: TORCH=1 python -m pytest -s -v
+ run: LAZY=0 TORCH=1 python -m pytest -s -v
- name: Run Pytest (lazy)
run: LAZY=1 TORCH=1 python -m pytest -s -v
@@ -97,7 +97,7 @@ jobs:
- name: Install Dependencies
run: pip install -e '.[gpu,testing]'
- name: Run Pytest
- run: GPU=1 python -m pytest -s -v
+ run: LAZY=0 GPU=1 python -m pytest -s -v
- name: Run Pytest (lazy)
run: LAZY=1 GPU=1 python -m pytest -s -v

MODIFY
examples/benchmark_train_efficientnet.py
examples/benchmark_train_efficientnet.py
@@ -25,6 +25,7 @@ if __name__ == "__main__":
Tensor.training = True
for i in trange(CNT):
+ cpy = time.monotonic()
x_train = Tensor.randn(BS, 3, 224, 224, requires_grad=False).realize()
y_train = Tensor.randn(BS, 1000, requires_grad=False).realize()
@@ -44,7 +45,7 @@ if __name__ == "__main__":
loss = loss.detach().cpu().data[0]
cl = time.monotonic()
- print(f"{(cl-st)*1000.0:7.2f} ms run, {(mt-st)*1000.0:7.2f} ms build, {(et-mt)*1000.0:7.2f} ms realize, {(cl-et)*1000.0:7.2f} ms CL, {loss:7.2f} loss, {tensors_allocated():4d} tensors, {mem_used/1e9:.2f} GB used")
+ print(f"{(st-cpy)*1000.0:7.2f} ms cpy, {(cl-st)*1000.0:7.2f} ms run, {(mt-st)*1000.0:7.2f} ms build, {(et-mt)*1000.0:7.2f} ms realize, {(cl-et)*1000.0:7.2f} ms CL, {loss:7.2f} loss, {tensors_allocated():4d} tensors, {mem_used/1e9:.2f} GB used")

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -3,17 +3,18 @@ import os, functools
import numpy as np
import pyopencl as cl # type: ignore
from collections import defaultdict
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Dict
from tinygrad.helpers import prod, ConvArgs
from tinygrad.ops import DEBUG, UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps
from tinygrad.shapetracker import ShapeTracker, View, strides_for_shape
-CLCACHE = int(os.getenv("CLCACHE", "0"))
+CLCACHE = int(os.getenv("CLCACHE", "1"))
class CLBuffer:
def __init__(self, size):
if len(CL.BUFFER_CACHE[size]) > 0: self.cl = CL.BUFFER_CACHE[size].pop()
else:
CL.mem_used += size
+ # TODO: on GPU OOM, clear the cache
self.cl = cl.Buffer(CL().cl_ctx, cl.mem_flags.READ_WRITE, size)
def __del__(self):
@@ -21,7 +22,8 @@ class CLBuffer:
else: CL.mem_used -= self.cl.size
class CL:
- CACHE, BUFFER_CACHE, kernel_count, mem_used = None, defaultdict(list), 0, 0
+ CACHE, kernel_count, mem_used = None, 0, 0
+ BUFFER_CACHE : Dict[int, List[cl.Buffer]] = defaultdict(list)
cl_ctx : Optional[cl.Context] = None
cl_queue : Optional[cl.CommandQueue] = None
def __init__(self):
@@ -30,6 +32,7 @@ class CL:
if len(devices) == 0: # settle for CPU
devices = sum([x.get_devices(device_type=cl.device_type.CPU) for x in cl.get_platforms()], [])
CL.cl_ctx = cl.Context(devices=[devices[int(os.getenv("CL_DEVICE", "0"))]])
+ if len(devices) > 1: print(f"using {CL.cl_ctx.devices}")
CL.cl_queue = cl.CommandQueue(self.cl_ctx) # this is an in-order command queue
@staticmethod
@@ -63,7 +66,7 @@ class GPUBuffer:
def __init__(self, shape, hostbuf:Optional[GPUBuffer]=None, backing:Optional[np.ndarray]=None):
self.st = ShapeTracker(shape)
self.shape = self.st.shape
- self._buf : CLBuffer = hostbuf._buf if hostbuf is not None else None
+ self._buf : Optional[CLBuffer] = hostbuf._buf if hostbuf is not None else None
self._base_shape : Tuple[int, ...] = hostbuf._base_shape if hostbuf is not None else self.shape
self._backing : Optional[np.ndarray] = hostbuf._backing if hostbuf is not None else backing
# early copy in for large buffers

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -148,7 +148,7 @@ class LazyOp(NamedTuple):
def get_lazybuffers(op:LazyOp) -> List[LazyBuffer]: return functools.reduce(operator.add, [get_lazybuffers(x) if isinstance(x, LazyOp) else [x] for x in op.src], [])
def get_lazyops(op:LazyOp) -> List[LazyOp]: return functools.reduce(operator.add, [get_lazyops(x) for x in op.src if isinstance(x, LazyOp)], [op])
-LAZY = int(os.getenv("LAZY", "0"))
+LAZY = int(os.getenv("LAZY", "1"))
class LazyBuffer:
def __init__(self, device, shape:Union[ShapeTracker, Tuple[int, ...]], optype:OpType, op:LazyOp):

LAZY and CLCACHE are defaults