batch from lazy branch
MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
@@ -5,10 +5,9 @@ def prod(x): return math.prod(x)
def reduce_shape(shape, axis):
return [1 if i in axis else shape[i] for i in range(len(shape))]
+conv_args = namedtuple('conv_args', ['H', 'W', 'groups', 'rcout', 'cin', 'oy', 'ox', 'iy', 'ix', 'ys', 'xs', 'bs', 'cout', 'py', 'px', 'dy', 'dx', 'out_shape'])
def get_conv_args(x_shape, w_shape, stride=1, groups=1, padding=0, dilation=1):
# TODO: https://docs.nvidia.com/deeplearning/performance/dl-performance-convolutional/index.html#tensor-layout
- conv_args = namedtuple('conv_args',
- ['H', 'W', 'groups', 'rcout', 'cin', 'oy', 'ox', 'iy', 'ix', 'ys', 'xs', 'bs', 'cout', 'py', 'px', 'dy','dx', 'out_shape'])
cout,cin,H,W = w_shape
ys,xs = (stride, stride) if isinstance(stride, int) else stride
py,px = (padding, padding) if isinstance(padding, int) else padding

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -6,6 +6,7 @@ from tinygrad.helpers import prod
from tinygrad.llops.ops_cpu import unary_op
from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps
from tinygrad.shapetracker import ShapeTracker, View, strides_for_shape
+from collections import defaultdict
cl_ctx, cl_queue = None, None
def get_cl_ctx(): return cl_ctx
@@ -25,22 +26,18 @@ class GPUBuffer:
def __init__(self, shape, hostbuf=None):
require_init_gpu()
self.st = ShapeTracker(shape)
- self.cl = hostbuf.cl if isinstance(hostbuf, GPUBuffer) else cl.Buffer(cl_ctx, cl.mem_flags.READ_WRITE, 4*roundup(prod(self.shape))) # padding
- if hostbuf is not None and not isinstance(hostbuf, GPUBuffer):
- cl.enqueue_copy(cl_queue, self.cl, hostbuf.astype(np.float32).ravel(), is_blocking=False)
-
- @property
- def shape(self): return self.st.shape
-
- @property
- def dtype(self): return np.float32
+ self.shape = self.st.shape
+ self.cl = hostbuf.cl if hostbuf is not None else cl.Buffer(cl_ctx, cl.mem_flags.READ_WRITE, 4*roundup(prod(self.shape))) # padding
def __repr__(self):
return f"<GPUBuffer with shape {self.shape!r}>"
@staticmethod
def fromCPU(x):
- return GPUBuffer(x.shape, x.view(np.ndarray))
+ ret = GPUBuffer(x.shape)
+ # TODO: this is blocking even though we told it not to
+ cl.enqueue_copy(cl_queue, ret.cl, x.view(np.ndarray).astype(np.float32).ravel(), is_blocking=False)
+ return ret
def toCPU(self):
data = np.empty(self.shape, dtype=np.float32)
@@ -49,38 +46,110 @@ class GPUBuffer:
class CLProgram:
def __init__(self, name, prg, options, argdtypes):
+ self.name = name
self.built = cl.Program(cl_ctx, prg).build(options=options)
self.clprg = self.built.__getattr__(name)
if argdtypes is not None: self.clprg.set_scalar_arg_dtypes(argdtypes)
- def __call__(self, *args): self.clprg(cl_queue, *args)
+ def __call__(self, *args):
+ #print(f"running {self.name} with {args[0]} count {len(args)-2}")
+ self.clprg(cl_queue, *args)
@functools.lru_cache(maxsize=None)
def clbuild(name, prg, options=tuple(), argdtypes=None):
- #print("cache miss")
- #print(prg)
+ #print("cache miss", prg[0:100])
return CLProgram(name, prg, options, argdtypes)
-code_for_op = {
- UnaryOps.NOOP: "A", UnaryOps.RELU: "max(A, (float)0.)", UnaryOps.EXP: "exp(A)", UnaryOps.LOG: "log(A)", UnaryOps.NEG: "-A", UnaryOps.SIGN: "sign(A)",
- BinaryOps.ADD: "A+B", BinaryOps.SUB: "A-B", BinaryOps.MUL: "A*B", BinaryOps.DIV: "B/A", BinaryOps.POW: "pow(A,B)", BinaryOps.CMPEQ: "(A==B)"
-}
-
-def contiguous_view(x:GPUBuffer, name:str):
+def contiguous_view(x:GPUBuffer, name:str) -> str:
return f"inline float get_{name}(__global const float *x, int gid) {{ int valid = 1; int idx = gid; {x.st.expr().replace('//', '/')}; return valid ? x[idx] : 0.0;}}"
-def elementwise_op(bufs: List[Tuple[str, GPUBuffer]], code):
- assert all(buf.shape == bufs[0][1].shape for _, buf in bufs)
+def elementwise_op_compile(bufs: List[Tuple[str, GPUBuffer]], code:str) -> str:
+ return '\n'.join([contiguous_view(buf, name) for name, buf in bufs])+ \
+ "inline float _ewop("+','.join(["int gid", "float acc"]+[f"__global const float *{name}_g" for name, _ in bufs])+") {"+ \
+ '\n'.join([f"float {name} = get_{name}({name}_g, gid);" for name, _ in bufs])+ \
+ f"return {code}; }}"
+
+def _processing_op(bufs: List[Tuple[str, GPUBuffer]]=[], code:str="acc", C=None):
+ if C is not None:
+ ret = GPUBuffer(C.out_shape)
+ ints = ''.join(f"int {x} = {getattr(C, x)};" for x in ["H", "W", "ys", "xs", "dx", "dy", "px", "py"] + (["cin"] if C.cin == 1 else []))
+ params = [(f"int {x}", getattr(C, x)) for x in ["groups", "rcout", "oy", "ox", "iy", "ix"] + (["cin"] if C.cin > 1 else [])]
+ options = []
+ if C.px == 0 and C.py == 0: options.append("-DALLVALID")
+ if C.oy == 1 and C.ox == 1: options.append("-DONEBYONE")
+ global_size = [C.bs*C.cout, C.oy, C.ox]
+ ewbufs = bufs[2:]
+ else:
ret = GPUBuffer(bufs[0][1].shape)
- ewop = clbuild("ewop", '\n'.join([contiguous_view(buf, name) for name, buf in bufs])+
- "__kernel void ewop(__global float *res_g, "+','.join([f"__global const float *{name}_g" for name, _ in bufs])+") {"+
- "int gid = get_global_id(0);"+
- '\n'.join([f"float {name} = get_{name}({name}_g, gid);" for name, _ in bufs])+
- f"res_g[gid] = {code}; }}")
- ewop([prod(ret.shape)], None, ret.cl, *[buf.cl for _, buf in bufs])
+ ints = ''
+ params = []
+ options = ["-DNOCONV"]
+ global_size = [prod(ret.shape), 1, 1]
+ ewbufs = bufs
+
+ conv_params = ["__global float* restrict output"] + \
+ [f"__global const float *{name}_g" for name, _ in bufs] + \
+ [x[0] for x in params]
+ conv_prg = clbuild("conv", elementwise_op_compile(ewbufs, code)+"""
+ __kernel void conv("""+','.join(conv_params)+""") {
+ float acc = 0.0;
+ int gid = get_global_id(0);
+ """+ints+"""
+
+#ifndef NOCONV
+ int B = get_global_id(0)/(groups*rcout); // range 0-bs
+ int g = (get_global_id(0)/rcout)%groups;
+ int c = get_global_id(0) % rcout;
+
+#ifdef ONEBYONE
+ int Y = 0;
+ int X = 0;
+#else
+ int Y = get_global_id(1); // range 0-oy
+ int X = get_global_id(2); // range 0-ox
+ gid = gid*oy*ox + Y*ox + X;
+#endif
+
+ int IY = Y*ys;
+ int IX = X*xs;
+
+ for (int ci = 0; ci < cin; ci++) {
+ for (int y = 0; y < H; y++) { for (int x = 0; x < W; x++) {
+ int idx_y = y*dy + IY - py;
+ int idx_x = x*dx + IX - px;
+#ifdef ALLVALID
+ acc += input_g[B*groups*cin*iy*ix + g*cin*iy*ix + ci*iy*ix + idx_y*ix + idx_x] * \
+ weight_g[g*rcout*cin*H*W + c*cin*H*W + ci*H*W + y*W + x];
+#else
+ int valid = (idx_y >= 0 && idx_y < iy && idx_x >= 0 && idx_x < ix);
+ acc += valid * input_g[B*groups*cin*iy*ix + g*cin*iy*ix + ci*iy*ix + clamp(idx_y, 0, iy-1)*ix + clamp(idx_x, 0, ix-1)] * \
+ weight_g[g*rcout*cin*H*W + c*cin*H*W + ci*H*W + y*W + x];
+#endif
+ } }
+ }
+#endif
+
+ output[gid] = _ewop("""+','.join(["gid", "acc"]+[f"{name}_g" for name, _ in ewbufs])+""");
+ }""", options=tuple(options),
+ argdtypes=tuple([None]*(1+len(bufs)) + [np.int32]*len(params)))
+ conv_prg(global_size, None, ret.cl, *[buf.cl for _, buf in bufs], *[x[1] for x in params])
+ return ret
+
+def movement_op(op, x, arg):
+ ret = GPUBuffer(x.st, x)
+ ret.shape = ret.st.movement_op(op, arg).shape
return ret
-def unary_op(op, x): return elementwise_op([("A", x)], code_for_op[op])
-def binary_op(op, x, y): return elementwise_op([("A", x), ("B", y)], code_for_op[op])
+def processing_op(op,x,w,C):
+ assert op == ProcessingOps.CONV, f"{op} isn't supported"
+ return _processing_op([("input", contiguous(x)), ("weight", contiguous(w))], "acc", C)
+
+code_for_op = {
+ UnaryOps.NOOP: "(A)", UnaryOps.RELU: "max(A, (float)0.)", UnaryOps.EXP: "exp(A)", UnaryOps.LOG: "log(A)", UnaryOps.NEG: "(-(A))", UnaryOps.SIGN: "sign(A)",
+ BinaryOps.ADD: "(A+B)", BinaryOps.SUB: "(A-B)", BinaryOps.MUL: "(A*B)", BinaryOps.DIV: "(B/A)", BinaryOps.POW: "pow(A,B)", BinaryOps.CMPEQ: "(A==B)",
+}
+
+def unary_op(op, x): return _processing_op([("A", x)], code_for_op[op])
+def binary_op(op, x, y): return _processing_op([("A", x), ("B", y)], code_for_op[op])
def contiguous(x:GPUBuffer): return x if x.st.contiguous else unary_op(UnaryOps.NOOP, x)
def reduce_op(op, inp, new_shape):
@@ -116,50 +185,3 @@ def reduce_op(op, inp, new_shape):
}"""
clbuild("reduce", prg)([prod(ret.shape)], None, inp.cl, ret.cl)
return ret
-
-def movement_op(op, x, arg):
- ret = GPUBuffer(x.st, x)
- ret.st.movement_op(op, arg)
- return ret
-
-def processing_op(op,x,w,C):
- ret = GPUBuffer((C.bs, C.cout, C.oy, C.ox))
- assert op == ProcessingOps.CONV, f"{op} isn't supported"
- ints = ''.join(f"int {x} = {getattr(C, x)};" for x in ["H", "W", "cin", "ys", "xs", "dx", "dy", "px", "py"])
- params = [(f"int {x}", getattr(C, x)) for x in ["groups", "rcout", "oy", "ox", "iy", "ix"]]
- conv_prg = clbuild("conv", """
- __kernel void conv(__global const float* restrict input, __global const float* restrict weight, __global float* restrict output,
- """+','.join([x[0] for x in params])+""") {
- """+ints+"""
- int B = get_global_id(0)/(groups*rcout); // range 0-bs
- int g = (get_global_id(0)/rcout)%groups;
- int c = get_global_id(0) % rcout;
-
- int Y = get_global_id(1); // range 0-oy
- int X = get_global_id(2); // range 0-ox
- int IY = Y*ys;
- int IX = X*xs;
-
- int gid = get_global_id(0)*oy*ox + Y*ox + X;
-
- float acc = 0.0;
- for (int ci = 0; ci < cin; ci++) {
- for (int y = 0; y < H; y++) { for (int x = 0; x < W; x++) {
- int idx_y = y*dy + IY - py;
- int idx_x = x*dx + IX - px;
-#ifdef ALLVALID
- acc += input[B*groups*cin*iy*ix + g*cin*iy*ix + ci*iy*ix + idx_y*ix + idx_x] * \
- weight[g*rcout*cin*H*W + c*cin*H*W + ci*H*W + y*W + x];
-#else
- int valid = (idx_y >= 0 && idx_y < iy && idx_x >= 0 && idx_x < ix);
- acc += valid * input[B*groups*cin*iy*ix + g*cin*iy*ix + ci*iy*ix + clamp(idx_y, 0, iy-1)*ix + clamp(idx_x, 0, ix-1)] * \
- weight[g*rcout*cin*H*W + c*cin*H*W + ci*H*W + y*W + x];
-#endif
- } }
- }
- output[gid] = acc;
- }""",
- options=tuple(["-DALLVALID"]) if C.px == 0 and C.py == 0 else tuple(),
- argdtypes=tuple([None, None, None] + [np.int32]*len(params)))
- conv_prg([C.bs*C.cout, C.oy, C.ox], None, contiguous(x).cl, contiguous(w).cl, ret.cl, *[x[1] for x in params])
- return ret

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -10,7 +10,7 @@ LoadOps = Enum("LoadOps", ["FROMCPU"])
from tinygrad.shapetracker import ShapeTracker
import os
-DEBUG = int(os.getenv("PRINT_LLOPS", "0"))
+DEBUG = int(os.getenv("DEBUG", "0"))
GRAPH = int(os.getenv("GRAPH", "0"))
from collections import defaultdict
cnts = defaultdict(int)
@@ -27,7 +27,7 @@ if GRAPH:
atexit.register(save_graph_exit)
global_num_max = 0
-def log_op(optype, op, ret, inp, dashed=False):
+def log_op(optype, op, ret, inp):
cnts[optype] += 1
if DEBUG: print(f"{op} : {', '.join([str(x.shape) for x in inp])} -> {ret.shape}")
if GRAPH:
@@ -42,8 +42,9 @@ def log_op(optype, op, ret, inp, dashed=False):
for x in inp:
if not isinstance(op, list): op = [op]
- sop = '.'.join([str(y).split(".")[1] for y in op][::-1])
- G.add_edge(nm(x), nm(ret), label=sop, color='#808080' if dashed else '', style='dashed' if dashed else '')
+ #sop = '.'.join([str(y).split(".")[1] for y in op][::-1])
+ sop = str(len(op))
+ G.add_edge(nm(x), nm(ret), label=sop)
if 'label' not in G.nodes[nm(x)]: G.nodes[nm(x)]['label'] = str(x.shape)
if nm(ret) not in G.nodes: G.add_node(nm(ret))
st = getattr(ret, "st", None)
@@ -79,7 +80,8 @@ class Ops:
ret = ctx.op.movement_op(op, x, tuple(arg))
if 'LAZY' not in ctx.device: log_op(MovementOps, op, ret, [x])
assert isinstance(ret, ctx.buffer)
- assert ret.shape == ShapeTracker(x.shape).movement_op(op, arg).shape
+ # this check is slow
+ #assert ret.shape == ShapeTracker(x.shape).movement_op(op, arg).shape
return ret
def processing_op(ctx, op:ProcessingOps, x, y, C):

MODIFY
tinygrad/shapetracker.py
tinygrad/shapetracker.py
# ShapeTracker allows movement operations to a buffer that don't require a copy to be made.
+from __future__ import annotations
+import functools
+from typing import Tuple, Union, List
from tinygrad.helpers import prod
-from functools import cached_property, reduce
-from itertools import chain
def divmodidx(acc, d, mod=True):
lr = f"(idx//{acc})" if acc != 1 else "idx"
return f"({lr}%{d})" if mod else lr # don't mod the top shape dimension
-class View:
- def __init__(self, shape, strides, offset:int=0):
- self.shape, self.strides, self.offset = tuple(shape), tuple(strides), offset
+@functools.lru_cache(maxsize=None)
+def to_shape_strides(shape:Tuple[int], strides:Tuple[int]) -> List[Tuple[int, int]]:
assert len(shape) == len(strides)
-
- self.shape_strides = [(shape[0], strides[0])]
+ ret = [(shape[0], strides[0])]
for i in range(1, len(shape)):
- if (strides[i] != 0 and self.shape_strides[-1][1]//strides[i] == shape[i]) or (strides[i] == 0 and self.shape_strides[-1][1] == 0):
- self.shape_strides[-1] = (self.shape_strides[-1][0] * shape[i], strides[i])
+ if (strides[i] != 0 and ret[-1][1]//strides[i] == shape[i]) or (strides[i] == 0 and ret[-1][1] == 0):
+ ret[-1] = (ret[-1][0] * shape[i], strides[i])
else:
- self.shape_strides.append((shape[i], strides[i]))
+ ret.append((shape[i], strides[i]))
+ return ret
+
+class View:
+ def __init__(self, shape, strides, offset:int=0):
+ self.shape, self.strides, self.offset = tuple(shape), tuple(strides), offset
+ self.shape_strides = to_shape_strides(self.shape, self.strides)
def __repr__(self): return f"View<{self.shape}, {self.strides}, {self.offset}>"
- @cached_property
+ @functools.cached_property
+ def contiguous(self):
+ return self.offset == 0 and all(s1 == s2 or s == 1 for s,s1,s2 in zip(self.shape, self.strides, strides_for_shape(self.shape)))
+
+ @functools.cached_property
def expr(self):
ret = [f"{self.offset}"] if self.offset != 0 else []
acc = 1
@@ -46,25 +55,27 @@ class ZeroView:
acc *= self.shape[0]
self.expr = 'valid=' + ' && '.join(expr)
+@functools.lru_cache(maxsize=None)
def strides_for_shape(shape):
strides = [1]
for d in shape[::-1][:-1]:
strides = [d*strides[0]] + strides
return tuple(strides)
-class ShapeTracker:
- def __init__(self, shape, strides=None):
- if isinstance(shape, ShapeTracker):
- self.views = shape.views[:]
- else:
+@functools.lru_cache(maxsize=None)
+def view_from_shape(shape:Tuple):
if len(shape) == 0: shape = (1,)
assert all([isinstance(x, int) for x in shape])
- self.views = [View(tuple(shape), strides_for_shape(shape) if strides == None else strides)]
+ return View(tuple(shape), strides_for_shape(shape))
+
+class ShapeTracker:
+ def __init__(self, shape:Union[ShapeTracker, Tuple[int]]):
+ if isinstance(shape, ShapeTracker): self.views = shape.views[:]
+ else: self.views = [view_from_shape(shape)]
@property
def contiguous(self):
- if len(self.views) > 1 or self.offset != 0: return False
- return all(s1 == s2 or s == 1 for s,s1,s2 in zip(self.shape, self.strides, strides_for_shape(self.shape)))
+ return len(self.views) == 1 and self.views[-1].contiguous
@property
def shape(self): return self.views[-1].shape

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -3,31 +3,7 @@ import os, atexit, time, inspect, functools, importlib
from collections import defaultdict
import numpy as np
from tinygrad.helpers import prod
-
-# **** profiler ****
-
-DEBUG = os.getenv("DEBUG", None) is not None
-if DEBUG:
- debug_counts, debug_times = defaultdict(int), defaultdict(float)
- def print_debug_exit():
- for name, _ in sorted(debug_times.items(), key=lambda x: -x[1]):
- print(f"{name:>20} : {debug_counts[name]:>6} {debug_times[name]:>10.2f} ms")
- atexit.register(print_debug_exit)
-
-global_num_max = 0
-class ProfileOp:
- def __init__(self, ctx, name, x, backward=False):
- self.ctx, self.name, self.x, self.output, self.backward = ctx, f"back_{name}" if backward else name, x, None, backward
- def __enter__(self):
- if DEBUG: self.st = time.time()
- return self
- def __exit__(self, *junk):
- if DEBUG:
- self.output[0].data.toCPU()
- et = (time.time()-self.st)*1000.
- debug_counts[self.name] += 1
- debug_times[self.name] += et
- print(f"{self.name:>20} : {et:>7.2f} ms {str([y.shape for y in self.x]):>40} -> {str([y.shape for y in self.output])}")
+from typing import List
# **** enumerate supported devices ****
@@ -65,6 +41,10 @@ class Tensor:
def __repr__(self):
return f"<Tensor {self.data!r} with grad {(self.grad.data if self.grad else None)!r}>"
+ def realize(self):
+ if getattr(self.data, 'realize', None) is not None:
+ self.data.realize()
+
def assign(self, x):
if not isinstance(x, Tensor):
x = Tensor(x)
@@ -78,7 +58,7 @@ class Tensor:
@staticmethod
def _get_data_dtype(data):
- return data.getdtype() if getattr(data, 'getdtype', None) else data.dtype
+ return data.getdtype() if getattr(data, 'getdtype', None) else (data.dtype if getattr(data, 'dtype', None) else np.float32)
@property
def dtype(self):
@@ -132,11 +112,9 @@ class Tensor:
if not any(x.requires_grad for x in t0._ctx.parents):
continue
assert (t0.grad is not None)
- with ProfileOp(t0._ctx, t0._ctx.__class__.__name__, [t0.grad], backward=True) as po:
grads = t0._ctx.backward(t0._ctx, t0.grad.data)
grads = [Tensor(g, device=self.device, requires_grad=False) if g is not None else None
for g in ([grads] if len(t0._ctx.parents) == 1 else grads)]
- po.output = [x for x in grads if x is not None] # backward can return None if no required gradient, don't profile it
for t, g in zip(t0._ctx.parents, grads):
if g is not None and t.requires_grad:
assert g.shape == t.shape, \
@@ -147,10 +125,11 @@ class Tensor:
@staticmethod
def _move_data(data, device):
- if isinstance(data, list):
- data = np.array(data, dtype=np.float32)
if isinstance(data, Device.buffers[device]):
return data
+ if isinstance(data, list):
+ # TODO: don't use np.array here, support Tensor creation direct to device
+ data = np.array(data, dtype=np.float32)
if isinstance(data, np.ndarray):
data = data.view(Device.buffers[Device.CPU])
@@ -354,14 +333,14 @@ class Tensor:
@staticmethod
def broadcasted(fxn, x, y):
tt = [arg for arg in [x,y] if isinstance(arg, Tensor)][0] # this is the prototype tensor
- if not isinstance(x, Tensor): x = Tensor(np.array([x], dtype=tt.dtype), device=tt.device, requires_grad=False)
- if not isinstance(y, Tensor): y = Tensor(np.array([y], dtype=tt.dtype), device=tt.device, requires_grad=False)
+ if not isinstance(x, Tensor): x = Tensor([x], device=tt.device, requires_grad=False)
+ if not isinstance(y, Tensor): y = Tensor([y], device=tt.device, requires_grad=False)
n_dims = max(len(x.shape), len(y.shape))
if len(x.shape) != n_dims: x = x.reshape(list(x.shape) + [1]*(n_dims-len(x.shape)))
if len(y.shape) != n_dims: y = y.reshape(list(y.shape) + [1]*(n_dims-len(y.shape)))
- shape_ret = tuple([int(x) for x in np.maximum(x.shape, y.shape)])
+ shape_ret = tuple([max(sx, sy) for sx,sy in zip(x.shape, y.shape)])
if x.shape != shape_ret: x = x.expand(shape_ret)
if y.shape != shape_ret: y = y.expand(shape_ret)
return fxn(x, y)
@@ -414,14 +393,11 @@ class Function(Ops):
self.saved_tensors.extend(x)
@classmethod
- def apply(cls, *x, **kwargs):
- assert all([isinstance(arg, Tensor) for arg in x])
+ def apply(cls, *x:List[Tensor], **kwargs):
ctx = cls(x[0].device, *x)
- with ProfileOp(ctx, ctx.__class__.__name__, x) as po:
ret = Tensor(cls.forward(ctx, *[t.data for t in x], **kwargs),
device=ctx.device, requires_grad=ctx.requires_grad)
- po.output = [ret]
- if ret.requires_grad:
+ if ctx.requires_grad:
ret._ctx = ctx # used by autograd engine
return ret

batch from lazy branch (#341)