still over line count, maybe test pass
MODIFY
test/test_llops.py
test/test_llops.py
@@ -4,6 +4,7 @@ import unittest
from tqdm import trange
import numpy as np
from tinygrad.tensor import Device
+from tinygrad.helpers import UnaryOps, BinaryOps, ReduceOps
if Device.DEFAULT == Device.GPU:
from tinygrad.llops.opencl import GPUBuffer, sync, unary_op, binary_op, reduce_op
@@ -20,7 +21,7 @@ class TestBenchmarkCL(unittest.TestCase):
shape = (1024,1024)
buf = GPUBuffer(shape, hostbuf=np.ones(shape, dtype=np.float32))
def fxn():
- unary_op('1.01f*a', buf, buf)
+ unary_op(UnaryOps.NEG, buf, buf)
its_sec = timeit(fxn, 100000, done=sync)
print(f"unary op (no sync) {its_sec:.2f} its/sec")
#self.assertGreater(its_sec, 10000)
@@ -29,7 +30,7 @@ class TestBenchmarkCL(unittest.TestCase):
shape = (1,)
buf = GPUBuffer(shape, hostbuf=np.ones(shape, dtype=np.float32))
def fxn():
- unary_op('1.01f*a', buf, buf)
+ unary_op(UnaryOps.NEG, buf, buf)
sync()
its_sec = timeit(fxn, 1000)
print(f"unary op tiny {its_sec:.2f} its/sec")
@@ -39,7 +40,7 @@ class TestBenchmarkCL(unittest.TestCase):
shape = (1024,1024)
buf = GPUBuffer(shape, hostbuf=np.ones(shape, dtype=np.float32))
def fxn():
- unary_op('1.01f*a', buf, buf)
+ unary_op(UnaryOps.NEG, buf, buf)
sync()
its_sec = timeit(fxn, 1000)
print(f"unary op {its_sec:.2f} its/sec")
@@ -50,7 +51,7 @@ class TestBenchmarkCL(unittest.TestCase):
buf_a = GPUBuffer(shape, hostbuf=np.ones(shape, dtype=np.float32))
buf_b = GPUBuffer(shape, hostbuf=np.ones(shape, dtype=np.float32))
def fxn():
- binary_op('a+b', buf_a, buf_b, buf_a)
+ binary_op(BinaryOps.ADD, buf_a, buf_b, buf_a)
its_sec = timeit(fxn, 10000, done=sync)
print(f"binary op (no sync) {its_sec:.2f} its/sec")
#self.assertGreater(its_sec, 1000)
@@ -60,7 +61,7 @@ class TestBenchmarkCL(unittest.TestCase):
buf_a = GPUBuffer(shape, hostbuf=np.ones(shape, dtype=np.float32))
buf_b = GPUBuffer(shape, hostbuf=np.ones(shape, dtype=np.float32))
def fxn():
- binary_op('a+b', buf_a, buf_b, buf_a)
+ binary_op(BinaryOps.ADD, buf_a, buf_b, buf_a)
sync()
its_sec = timeit(fxn, 1000)
print(f"binary op {its_sec:.2f} its/sec")
@@ -71,7 +72,7 @@ class TestBenchmarkCL(unittest.TestCase):
buf = GPUBuffer(shape, hostbuf=np.ones(shape, dtype=np.float32))
buf_out = GPUBuffer((1024,1))
def fxn():
- reduce_op('out += a', buf, buf_out)
+ reduce_op(ReduceOps.SUM, buf, buf_out)
sync()
its_sec = timeit(fxn, 100)
print(f"reduce op {its_sec:.2f} its/sec")
@@ -82,7 +83,7 @@ class TestBenchmarkCL(unittest.TestCase):
buf = GPUBuffer(shape, hostbuf=np.ones(shape, dtype=np.float32))
buf_out = GPUBuffer((1,1))
def fxn():
- reduce_op('out += a', buf, buf_out)
+ reduce_op(ReduceOps.SUM, buf, buf_out)
sync()
its_sec = timeit(fxn, 10)
print(f"reduce op full {its_sec:.2f} its/sec")

MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
@@ -27,6 +27,7 @@ class UnaryOps(Enum):
EXP = 1
LOG = 2
NEG = 3
+ SIGN = 4
class BinaryOps(Enum):
ADD = 0
@@ -37,13 +38,9 @@ class BinaryOps(Enum):
A = 5
CMP = 6
# these are "complex" ops, and should probably be replaced
- EXPMUL = 7
- POW_D1 = 8
- POW_D2 = 9
- RELU_D = 10
+ POW_D1 = 7
+ POW_D2 = 8
class ReduceOps(Enum):
SUM = 0
MAX = 1
\ No newline at end of file
- # this is "complex"
- NEGSUM = 2
\ No newline at end of file

MODIFY
tinygrad/llops/opencl.py
tinygrad/llops/opencl.py
@@ -53,6 +53,7 @@ def unary_op(op, x, ret):
elif op == UnaryOps.EXP: code = 'exp(a)'
elif op == UnaryOps.LOG: code = 'log(a)'
elif op == UnaryOps.NEG: code = '-a'
+ elif op == UnaryOps.SIGN: code = 'sign(a)'
else: raise Exception(f"{op} isn't supported")
unop = clbuild("unop", """
__kernel void unop(__global const float4 *a_g, __global float4 *res_g) {
@@ -91,10 +92,8 @@ def binary_op(op, x, y, ret):
elif op == BinaryOps.POW: code = "pow(a,b)"
elif op == BinaryOps.A: code = "a"
elif op == BinaryOps.CMP: code = "1.0f*(a==b)"
- elif op == BinaryOps.EXPMUL: code = "exp(a) * b"
elif op == BinaryOps.POW_D1: code = 'b * pow(a, b-1.0f)'
elif op == BinaryOps.POW_D2: code = 'log(a) * pow(a, b)'
- elif op == BinaryOps.RELU_D: code = 'b * max(sign(a), (float)0.)'
else: raise Exception(f"{op} isn't supported")
shape_ret, dimlist, complist = binary_broadcast(x.shape, y.shape, True)
@@ -108,7 +107,6 @@ def binary_op(op, x, y, ret):
def reduce_op(op, inp, ret, start="0.0"):
if op == ReduceOps.SUM: code = "out += a"
elif op == ReduceOps.MAX: code = "out = max(a,out)"
- elif op == ReduceOps.NEGSUM: code = "out -= a"
else: raise Exception(f"{op} isn't supported")
# TODO: this is insanely slow
# NOTE: ret.shape can be (1,), it's mostly by luck that this works

MODIFY
tinygrad/ops/ops_gpu.py
tinygrad/ops/ops_gpu.py
@@ -18,15 +18,25 @@ class UnaryOp(Function):
class ReLU(UnaryOp):
fop = UnaryOps.RELU
- bop = BinaryOps.RELU_D
+
+ def backward(ctx, grad_output):
+ input, = ctx.saved_tensors
+ ret = Buffer(input.shape)
+ unary_op(UnaryOps.SIGN, input, ret)
+ unary_op(UnaryOps.RELU, ret, ret)
+ return binary_op(BinaryOps.MUL, ret, grad_output, ret)
class Log(UnaryOp):
fop = UnaryOps.LOG
bop = BinaryOps.DIV
class Exp(UnaryOp):
- fop = UnaryOps.EXP
- bop = BinaryOps.EXPMUL
+ def forward(ctx, input):
+ ret = unary_op(UnaryOps.EXP, input, Buffer(input.shape))
+ ctx.save_for_backward(ret) # we save the output here, not the input
+ return ret
+
+ bop = BinaryOps.MUL
# ************* reduce ops *************
@@ -59,8 +69,8 @@ class Max(Function):
# ************* binary ops *************
-def unbroadcast(out, in_sh, op=ReduceOps.SUM):
- return reduce_op(op, out, Buffer(in_sh))
+def unbroadcast(out, in_sh):
+ return reduce_op(ReduceOps.SUM, out, Buffer(in_sh))
class Add(Function):
def forward(ctx, x, y):
@@ -79,8 +89,9 @@ class Sub(Function):
def backward(ctx, grad_output):
shape_x, shape_y = ctx.saved_tensors
+ neg_grad_output = unary_op(UnaryOps.NEG, grad_output, Buffer(grad_output.shape))
return unbroadcast(grad_output, shape_x) if ctx.needs_input_grad[0] else None, \
- unbroadcast(grad_output, shape_y, ReduceOps.NEGSUM) if ctx.needs_input_grad[1] else None
+ unbroadcast(neg_grad_output, shape_y) if ctx.needs_input_grad[1] else None
class Mul(Function):
def forward(ctx, x, y):
@@ -96,11 +107,12 @@ class Mul(Function):
class Pow(Function):
def forward(ctx, x, y):
- ctx.save_for_backward(x, y)
- return binary_op(BinaryOps.POW, x, y, Buffer(binary_broadcast(x.shape, y.shape)))
+ ret = Buffer(binary_broadcast(x.shape, y.shape))
+ ctx.save_for_backward(x, y, ret)
+ return binary_op(BinaryOps.POW, x, y, ret)
def backward(ctx, grad_output):
- x,y = ctx.saved_tensors
+ x,y,ret = ctx.saved_tensors
tmp = Buffer(grad_output.shape)
grad_x = unbroadcast(binary_op(BinaryOps.MUL, grad_output, binary_op(BinaryOps.POW_D1, x, y, tmp), tmp), x.shape) if ctx.needs_input_grad[0] else None
grad_y = unbroadcast(binary_op(BinaryOps.MUL, grad_output, binary_op(BinaryOps.POW_D2, x, y, tmp), tmp), y.shape) if ctx.needs_input_grad[1] else None

still over line count, maybe test pass