stuff from lazy
MODIFY
test/test_shapetracker.py
test/test_shapetracker.py
@@ -36,10 +36,41 @@ class DumbShapeTracker:
# Tensor.zeros(2, 4).permute(1,0).reshape(2, 4)
# (d1*4 + d0%4), d1=x//4, d0=x%4 = ((x//4)*4) + (x%4)%4
+class TestSingleShapeTracker(unittest.TestCase):
+ def setUp(self):
+ self.st = ShapeTracker(7,4)
+
+ def test_reshape(self):
+ self.st.reshape(7,1,4)
+ assert self.st.contiguous
+
+ def test_permute(self):
+ self.st.permute(1,0)
+ assert not self.st.contiguous
+
+ def test_slice(self):
+ self.st.slice((1,2), (0,4))
+ assert not self.st.contiguous
+
+ def test_double_permute(self):
+ self.st.permute(1,0)
+ self.st.permute(1,0)
+ assert self.st.contiguous
+
+ def test_reshape_permute(self):
+ self.st.reshape(7,1,4)
+ self.st.permute(0,1,2)
+ assert self.st.contiguous
+
+ def test_reshape_permute_no(self):
+ self.st.reshape(4,7)
+ self.st.permute(1,0)
+ assert not self.st.contiguous
+
class TestShapeTracker(unittest.TestCase):
def setUp(self):
- self.st = ShapeTracker(2,4)
- self.dt = DumbShapeTracker(2,4)
+ self.st = ShapeTracker(7,4)
+ self.dt = DumbShapeTracker(7,4)
self.apply = lambda fxn: [fxn(x) for x in [self.st, self.dt]]
def tearDown(self):
@@ -54,7 +85,7 @@ class TestShapeTracker(unittest.TestCase):
def test_simple_split(self):
self.test_permute()
- self.apply(lambda x: x.reshape(8))
+ self.apply(lambda x: x.reshape(prod(self.st.shape)))
def test_reshape(self):
assert self.st.shape == self.dt.shape
@@ -98,6 +129,22 @@ class TestShapeTracker(unittest.TestCase):
def test_slice_1c2(self):
self.apply(lambda x: x.slice((1, 2), (1, 2)))
+ def test_double_permute(self):
+ self.apply(lambda x: x.permute(1, 0))
+ self.apply(lambda x: x.permute(1, 0))
+
+ def test_slice_permute(self):
+ self.apply(lambda x: x.slice((0, 2), (2, 4)))
+ self.apply(lambda x: x.permute(1, 0))
+
+ def test_slice_expand(self):
+ self.apply(lambda x: x.slice((0, 2), (3, 4)))
+ self.apply(lambda x: x.expand(2, 10))
+
+ def test_double_stride(self):
+ self.apply(lambda x: x.stride(1, 2))
+ self.apply(lambda x: x.stride(2, 1))
+
def test_stride(self): self.apply(lambda x: x.stride(2,1))
def test_stride_int(self): self.apply(lambda x: x.stride(1,2))
def test_stride_2(self): self.apply(lambda x: x.stride(2,2))

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -23,10 +23,9 @@ class GPUBuffer:
def __init__(self, shape, hostbuf=None):
require_init_gpu()
self.shape = tuple(shape)
- self.cl = cl.Buffer(cl_ctx, cl.mem_flags.READ_WRITE, 4*roundup(prod(self.shape))) # padding
- if hostbuf is not None:
- # TODO: this doesn't have to block
- cl.enqueue_copy(cl_queue, self.cl, hostbuf.astype(np.float32).ravel())
+ 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 dtype(self): return np.float32
@@ -128,15 +127,16 @@ def contiguous(x, st, ret=None):
return ret
def movement_op(op, x, arg=None):
- return contiguous(x, ShapeTracker(*x.shape).movement_op(op, arg))
+ st = ShapeTracker(*x.shape).movement_op(op, arg)
+ if st.contiguous: return GPUBuffer(st.shape, x)
+ else: return contiguous(x, st)
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"
conv_prg = clbuild("conv", """
- __kernel void conv(__global const float *input, __global const float *weight, __global float *output,
+ __kernel void conv(__global const float* restrict input, __global const float* restrict weight, __global float* restrict output,
int H, int W, int groups, int rcout, int cin, int oy, int ox, int iy, int ix, int ys, int xs, int bs, int dx, int dy, int px, int py) {
-
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;

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -5,24 +5,30 @@ BinaryOps = Enum("BinaryOps", ["ADD", "SUB", "MUL", "DIV", "POW", "CMPEQ"])
ReduceOps = Enum("ReduceOps", ["SUM", "MAX"])
MovementOps = Enum("MovementOps", ["RESHAPE", "PERMUTE", "SLICE", "EXPAND", "FLIP"])
ProcessingOps = Enum("ProcessingOps", ["CONV"])
+LoadOps = Enum("LoadOps", ["FROMCPU"])
from tinygrad.shapetracker import ShapeTracker
import os
DEBUG = int(os.getenv("PRINT_LLOPS", "0"))
GRAPH = int(os.getenv("GRAPH", "0"))
+from collections import defaultdict
+cnts = defaultdict(int)
if GRAPH:
import atexit
import networkx as nx
G = nx.DiGraph()
def save_graph_exit():
+ for k,v in cnts.items():
+ print(k, v)
print("saving", G)
nx.drawing.nx_pydot.write_dot(G, '/tmp/net.dot')
os.system('dot -Tsvg /tmp/net.dot -o /tmp/net.svg')
atexit.register(save_graph_exit)
global_num_max = 0
-def log_op(op, ret, inp):
+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:
def nm(x):
@@ -32,27 +38,28 @@ def log_op(op, ret, inp):
global_num_max += 1
return f"<<< {x.global_num} >>>"
- top,sop = str(op).split(".")
- top_colors = {"UnaryOps": "#c0c0c0", "ReduceOps": "#8080ff", "BinaryOps": "#c0c0c0", "MovementOps": "#80ff80", "ProcessingOps": "#ff8080"}
+ _,sop = str(op).split(".")
+ top_colors = {LoadOps: "#ffffff", UnaryOps: "#c0c0c0", ReduceOps: "#8080ff", BinaryOps: "#c0c0c0", MovementOps: "#80ff80", ProcessingOps: "#ff8080"}
for x in inp:
G.add_edge(nm(x), nm(ret), label=sop)
- G.nodes[nm(x)]['label'] = str(x.shape)
+ 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))
G.nodes[nm(ret)]['label'] = str(ret.shape)
- G.nodes[nm(ret)]['fillcolor'] = top_colors[top]
+ G.nodes[nm(ret)]['fillcolor'] = top_colors[optype]
G.nodes[nm(ret)]['style'] = 'filled'
class Ops:
def unary_op(ctx, op:UnaryOps, x):
ret = ctx.op.unary_op(op, x)
- log_op(op, ret, [x])
+ if 'LAZY' not in ctx.device: log_op(UnaryOps, op, ret, [x])
assert isinstance(ret, ctx.buffer)
assert ret.shape == x.shape
return ret
def reduce_op(ctx, op:ReduceOps, x, new_shape):
- ret = ctx.op.reduce_op(op, x, new_shape)
- log_op(op, ret, [x])
+ ret = ctx.op.reduce_op(op, x, tuple(new_shape))
+ if 'LAZY' not in ctx.device: log_op(ReduceOps, op, ret, [x])
assert isinstance(ret, ctx.buffer)
assert ret.shape == tuple(new_shape)
return ret
@@ -60,21 +67,21 @@ class Ops:
def binary_op(ctx, op:BinaryOps, x, y):
assert x.shape == y.shape
ret = ctx.op.binary_op(op, x, y)
- log_op(op, ret, [x, y])
+ if 'LAZY' not in ctx.device: log_op(BinaryOps, op, ret, [x, y])
assert isinstance(ret, ctx.buffer)
assert ret.shape == x.shape
return ret
- def movement_op(ctx, op:MovementOps, x, arg=None):
- ret = ctx.op.movement_op(op, x, arg)
- log_op(op, ret, [x])
+ def movement_op(ctx, op:MovementOps, x, arg):
+ 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
return ret
def processing_op(ctx, op:ProcessingOps, x, y, C):
ret = ctx.op.processing_op(op, x, y, C)
- log_op(op, ret, [x, y])
+ if 'LAZY' not in ctx.device: log_op(ProcessingOps, op, ret, [x, y])
assert isinstance(ret, ctx.buffer)
assert ret.shape == (C.bs, C.cout, C.oy, C.ox)
return ret
\ No newline at end of file

MODIFY
tinygrad/shapetracker.py
tinygrad/shapetracker.py
@@ -10,7 +10,7 @@ def divmodidx(acc, d, mod=True):
class View:
def __init__(self, shape, strides, offset:int=0):
assert len(shape) == len(strides)
- self.shape, self.offset = shape, offset
+ self.shape, self.strides, self.offset = tuple(shape), tuple(strides), offset
self.shape_strides = [(shape[0], strides[0])]
for i in range(1, len(shape)):
@@ -48,19 +48,27 @@ def strides_for_shape(shape):
strides = [1]
for d in shape[::-1][:-1]:
strides = [d*strides[0]] + strides
- return strides
+ return tuple(strides)
-# TODO: support "contiguous" property and "simple_strided" property and make them work
-# TODO: support simplification across views
class ShapeTracker:
def __init__(self, *shape, strides=None):
assert all([isinstance(x, int) for x in shape])
if len(shape) == 0: shape = (1,)
self.views = [View(shape, strides_for_shape(shape) if strides == None else strides)]
- self.contiguous = True
@property
- def shape(self): return tuple(self.views[-1].shape)
+ def contiguous(self):
+ if len(self.views) > 1: return False
+ return self.strides == strides_for_shape(self.shape) and self.offset == 0
+
+ @property
+ def shape(self): return self.views[-1].shape
+
+ @property
+ def strides(self): return self.views[-1].strides
+
+ @property
+ def offset(self): return self.views[-1].offset
def expr(self): return ';'.join([v.expr for v in self.views[::-1] if v.expr != 'idx=idx' and v.expr != 'valid=valid'])
def movement_op(self, op, arg): getattr(self, str(op).split(".")[1].lower())(*arg); return self
@@ -75,15 +83,14 @@ class ShapeTracker:
assert all([isinstance(x, int) for x in new_shape])
assert prod(self.shape) == prod(new_shape)
if self.shape == new_shape: return
- self.views.append(View(new_shape, strides_for_shape(new_shape)))
+ view = View(new_shape, strides_for_shape(new_shape))
+ if self.contiguous: self.views[-1] = view
+ else: self.views.append(view)
def permute(self, *axis):
assert all([isinstance(x, int) and x >= 0 and x < len(self.shape) for x in axis])
assert len(set(axis)) == len(axis) and len(axis) == len(self.shape)
- if tuple(range(len(axis))) == axis: return
- self.contiguous = False
- strides = strides_for_shape(self.shape)
- self.views.append(View([self.shape[a] for a in axis], [strides[a] for a in axis]))
+ self.views[-1] = View([self.shape[a] for a in axis], [self.strides[a] for a in axis], self.offset)
# TODO: this is a special case of slice with strides, remove it
# though it's nice that it can't change size
@@ -94,30 +101,24 @@ class ShapeTracker:
def slice(self, *arg):
assert len(arg) == len(self.shape)
- if all([(x,y) == (0,s) for s,(x,y) in zip(self.shape, arg)]): return
- self.contiguous = False
- strides = strides_for_shape(self.shape)
- offset = sum([strides[i]*x for i,(x,_) in enumerate(arg)])
+ offset = sum([self.strides[i]*x for i,(x,_) in enumerate(arg)])
zeroview = ZeroView(self.shape, arg)
- self.views.append(View([y-x for x,y in arg], strides, offset))
- if zeroview.expr != "valid=valid": self.views.append(zeroview)
+ self.views[-1] = View([y-x for x,y in arg], self.strides, self.offset+offset)
+ if zeroview.expr != "valid=valid":
+ # if we add a ZeroView, we add another (stock) view also for modding
+ self.views += [zeroview, View(self.shape, strides_for_shape(self.shape))]
def expand(self, *new_shape):
assert all([isinstance(x, int) for x in new_shape])
assert all([x == y or x == 1 for x,y in zip(self.shape, new_shape)])
- if self.shape == new_shape: return
- self.contiguous = False
- strides = [s if x == y else 0 for s,(x,y) in zip(strides_for_shape(self.shape), zip(self.shape, new_shape))]
- self.views.append(View(new_shape, strides))
+ strides = [s if x == y else 0 for s,(x,y) in zip(self.strides, zip(self.shape, new_shape))]
+ self.views[-1] = View(new_shape, strides, self.offset)
# TODO: combine with slice? this doesn't require a ZeroView, though slice shouldn't always either
def stride(self, *mul):
assert all([isinstance(x, int) for x in mul])
- if all([x==1 for x in mul]): return
- self.contiguous = False
- old_strides = strides_for_shape(self.shape)
- strides = [z*m for z,m in zip(old_strides, mul)]
+ strides = [z*m for z,m in zip(self.strides, mul)]
new_shape = [(s+(abs(m)-1))//abs(m) for s,m in zip(self.shape, mul)]
- offset = sum([(s-1)*z for s,z,m in zip(self.shape,old_strides,mul) if m < 0])
- self.views.append(View(new_shape, strides, offset))
+ offset = sum([(s-1)*z for s,z,m in zip(self.shape, self.strides, mul) if m < 0])
+ self.views[-1] = View(new_shape, strides, self.offset + offset)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -38,12 +38,12 @@ class Device:
buffers, llops = {}, {}
for i,op in imports.items():
name = op[len("ops_"):].upper()
- vars()[name] = i
- DEFAULT = i if os.environ.get(name, 0) == "1" else DEFAULT
+ vars()[name] = name
+ DEFAULT = name if os.environ.get(name, 0) == "1" else DEFAULT
try:
- llops[i] = importlib.import_module('tinygrad.llops.'+op)
+ llops[name] = importlib.import_module('tinygrad.llops.'+op)
def find_buffer(llo, name): return [cls for cname, cls in inspect.getmembers(llo, inspect.isclass) if (cname.upper() == name + "BUFFER")][0]
- buffers[i] = find_buffer(llops[i], name)
+ buffers[name] = find_buffer(llops[name], name)
except ImportError as e:
print(op, "not available", e)
DEFAULT = CPU if DEFAULT is None else DEFAULT

stuff from lazy