seperate STRIDED and EXPAND
MODIFY
accel/llvm/ops_llvm.py
accel/llvm/ops_llvm.py
@@ -21,7 +21,7 @@ def idx_deref(builder, buf, ptr, eidx):
idx = builder.add(builder.mul(eidx[1], int_const(eidx[2])), eidx[3], name="idx")
if DEBUG >= 1:
- print("viewcount:", len(buf.st.views), buf.st.expr(), ptr)
+ print("viewcount:", len(buf.st.views), buf.st.expr(), ptr, "on", buf.shape)
# TODO: unify this with expr in ShapeTracker
valid = None
for v in buf.st.views[::-1]:

MODIFY
test/test_ops.py
test/test_ops.py
@@ -31,7 +31,7 @@ def helper_test_op(shps, torch_fxn, tinygrad_fxn, atol=1e-6, rtol=1e-3, grad_ato
except Exception:
raise Exception(f"{s} failed shape {x.shape}")
- compare("forward pass", ret.cpu().data, out.detach().numpy(), atol=atol, rtol=rtol)
+ compare("forward pass", ret.numpy(), out.detach().numpy(), atol=atol, rtol=rtol)
torch_fbp, tinygrad_fbp = np.nan, np.nan
if not forward_only and not FORWARD_ONLY:
@@ -45,7 +45,7 @@ def helper_test_op(shps, torch_fxn, tinygrad_fxn, atol=1e-6, rtol=1e-3, grad_ato
tinygrad_fbp = time.monotonic() - st
for i, (t, tt) in enumerate(zip(ts, tst)):
- compare(f"backward pass tensor {i}", tt.cpu().grad.data, t.grad.detach().numpy(), atol=grad_atol, rtol=grad_rtol)
+ compare(f"backward pass tensor {i}", tt.grad.numpy(), t.grad.detach().numpy(), atol=grad_atol, rtol=grad_rtol)
print("\ntesting %40r torch/tinygrad fp: %.2f / %.2f ms bp: %.2f / %.2f ms " % (shps, torch_fp*1000, tinygrad_fp*1000, torch_fbp*1000, tinygrad_fbp*1000), end="")
@@ -104,6 +104,8 @@ class TestOps(unittest.TestCase):
helper_test_op([(2), (2,2)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-4)
def test_matmul(self):
helper_test_op([(65), (65,99)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-4)
+ def test_gemm(self):
+ helper_test_op([(256,256), (256,256)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-3)
def test_broadcastdot(self):
helper_test_op([(10,45,65), (65,45)], lambda x,y: x @ y, Tensor.dot, atol=1e-4)
def test_multidot(self):

MODIFY
tinygrad/lazy.py
tinygrad/lazy.py
@@ -228,8 +228,11 @@ class LazyBuffer:
# TODO: is there any way to replace strided with other movement ops?
x = x.movement_op(MovementOps.STRIDED, (
(C.bs, C.groups*C.cin*x.shape[2]*x.shape[3]), (C.groups, C.cin*x.shape[2]*x.shape[3]),
- (C.rcout, 0), (C.oy, C.sy*x.shape[3]), (C.ox, C.sx),
+ (1, 1), (C.oy, C.sy*x.shape[3]), (C.ox, C.sx),
(C.cin, x.shape[2]*x.shape[3]), (C.H, C.dy*x.shape[3]), (C.W, C.dx)))
+ #if C.H <= 3 and C.W <= 3: # max 9x the RAM overhead, this is im2col
+ # x = x.contiguous_op()
+ x = x.movement_op(MovementOps.EXPAND, (C.bs, C.groups, C.rcout, C.oy, C.ox, C.cin, C.H, C.W))
w = w.movement_op(MovementOps.RESHAPE, (1, C.groups, C.rcout, 1, 1, C.cin, C.H, C.W)) \
.movement_op(MovementOps.EXPAND, (C.bs, C.groups, C.rcout, C.oy, C.ox, C.cin, C.H, C.W))
return x.binary_op(BinaryOps.MUL, w).reduce_op(ReduceOps.SUM, (C.bs, C.groups, C.rcout, C.oy, C.ox, 1, 1, 1)) \

seperate STRIDED and EXPAND