cleanup, add flops tracking
MODIFY
accel/lazy/ops_lazy.py
accel/lazy/ops_lazy.py
@@ -92,7 +92,7 @@ class LazyBuffer:
@functools.lru_cache(maxsize=None if CACHE_LAZYBUFFERS else 0)
def movement_op(x:LazyBuffer, op:MovementOps, arg) -> LazyBuffer:
# TODO: SHUFFLE_SLICE_OPS is okay if it's a shrink
- if SHUFFLE_MOVEMENT_OPS and x.optype == BinaryOps and (SHUFFLE_SLICE_OPS or op != MovementOps.SLICE):
+ if SHUFFLE_MOVEMENT_OPS and x.optype == BinaryOps and x.realized is None and (SHUFFLE_SLICE_OPS or op != MovementOps.SLICE):
# if this MovementOp is being applied to a BinaryOp, apply the MovementOp to all the BinaryOp inputs instead
def replace_with_movement_op(y:Union[LazyOp, LazyBuffer]) -> LazyBuffer:
if isinstance(y, LazyBuffer): return y.movement_op(op, arg)
@@ -137,22 +137,20 @@ def elementwise_op(op, srcs:Tuple[LazyBuffer]) -> LazyBuffer:
if MERGE_ELEMENTWISE_INTO_CONV_OUTPUT:
cnt = sum([x.optype == ProcessingOps and x.realized is None for x in srcs])
+ if cnt > 0:
if cnt == 1:
srcs = [x.op if x.optype == ProcessingOps and x.realized is None else x for x in srcs]
- return LazyBuffer(out_shape, ProcessingOps, LazyOp(op, srcs))
elif cnt == 2:
- # have to confirm they are the same conv
- c1, c2 = [find_conv(x.op) for x in srcs]
- if c1.op == c1.op and c1.arg == c2.arg and tuple(c1.src) == tuple(c2.src):
+ c1, c2 = [find_conv(x.op) for x in srcs] # this returns a list of LazyOps
+ if c1 == c2: # NOTE: this compare relies on working caching for processing_ops
+ # they are the same conv, merge them
srcs = [x.op if x.optype == ProcessingOps and x.realized is None else x for x in srcs]
- return LazyBuffer(out_shape, ProcessingOps, LazyOp(op, srcs))
else:
+ # they are not the same conv, choose one
order = cmp(srcs[0], srcs[1])
- if order == -1:
- #if depends(srcs[0], srcs[1]):
+ if order == -1: #if depends(srcs[0], srcs[1]):
srcs = [srcs[0].op, srcs[1]]
- elif order == 1:
- #elif depends(srcs[1], srcs[0]):
+ elif order == 1: #elif depends(srcs[1], srcs[0]):
srcs = [srcs[0], srcs[1].op]
else:
# all three are okay

MODIFY
test/test_conv.py
test/test_conv.py
@@ -28,6 +28,13 @@ class TestConv(unittest.TestCase):
print(ret.numpy())
+ def test_two_binops_no_rerun(self):
+ x = Tensor.ones(1,12,128,256)
+ w = Tensor.ones(32,12,3,3)
+ out = x.conv2d(w, stride=(2,2), padding=(1,1))
+ out.relu().numpy(), (out-1).numpy()
+ # TODO: make this a real test
+
def test_first_three(self):
x = Tensor.ones(1,12,128,256)

MODIFY
tinygrad/ops.py
tinygrad/ops.py
from enum import Enum
+from tinygrad.helpers import prod
UnaryOps = Enum("UnaryOps", ["NOOP", "RELU", "EXP", "LOG", "NEG", "SIGN"])
BinaryOps = Enum("BinaryOps", ["ADD", "SUB", "MUL", "DIV", "POW", "CMPEQ"])
ReduceOps = Enum("ReduceOps", ["SUM", "MAX"])
@@ -10,13 +11,19 @@ DEBUG = int(os.getenv("DEBUG", "0"))
GRAPH = int(os.getenv("GRAPH", "0"))
from collections import defaultdict
cnts = defaultdict(int)
-if GRAPH:
+
import atexit
+if DEBUG:
+ def debug_exit():
+ for k,v in cnts.items():
+ print(k, v)
+ print(f"GFLOP: {Ops.flops*1e-9:.2f}")
+ atexit.register(debug_exit)
+
+if GRAPH:
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')
@@ -24,9 +31,9 @@ if GRAPH:
global_num_max = 0
def log_op(optype, op, ret, inp):
- if DEBUG: print(f"{op} : {', '.join([str(x.shape) for x in inp])} -> {ret.shape}")
- if GRAPH:
cnts[optype] += 1
+ if DEBUG >= 2: print(f"{op} : {', '.join([str(x.shape) for x in inp])} -> {ret.shape}")
+ if GRAPH:
def nm(x):
global global_num_max
if getattr(x, 'global_num', None) is None:
@@ -50,9 +57,12 @@ def log_op(optype, op, ret, inp):
G.nodes[nm(ret)]['style'] = 'filled, dashed' if non_contiguous else 'filled'
class Ops:
+ flops = 0
+
def unary_op(ctx, op:UnaryOps, x):
ret = x.unary_op(op)
if 'LAZY' not in ctx.device: log_op(UnaryOps, op, ret, [x])
+ Ops.flops += prod(x.shape)
assert isinstance(ret, ctx.buffer)
assert ret.shape == x.shape
return ret
@@ -68,6 +78,7 @@ class Ops:
assert x.shape == y.shape
ret = x.binary_op(op, y)
if 'LAZY' not in ctx.device: log_op(BinaryOps, op, ret, [x, y])
+ Ops.flops += prod(x.shape)*2
assert isinstance(ret, ctx.buffer)
assert ret.shape == x.shape
return ret
@@ -83,6 +94,7 @@ class Ops:
def processing_op(ctx, op:ProcessingOps, x, y, C):
ret = x.processing_op(op, y, C)
if 'LAZY' not in ctx.device: log_op(ProcessingOps, op, ret, [x, y])
+ Ops.flops += C.bs*C.cout*C.oy*C.ox*C.cin*C.H*C.W*2
assert isinstance(ret, ctx.buffer)
assert ret.shape == C.out_shape
return ret
\ No newline at end of file

cleanup, add flops tracking