Remove Matmul
MODIFY
README.md
README.md
@@ -116,14 +116,14 @@ hlops are syntactic sugar around mlops.
### mlops
-mlops are mid level ops, there's 14 of them. They understand memory allocation and derivatives
+mlops are mid level ops, there's 13 of them. They understand memory allocation and derivatives
```
Relu, Log, Exp # unary ops
Sum, Max # reduce ops (with axis argument)
Add, Sub, Mul, Pow # binary ops (with broadcasting)
Reshape, Transpose, Slice # movement ops
-Matmul, Conv2D(NCHW) # processing ops
+Conv2D(NCHW) # processing op (Matmul is also Conv2D)
```
You no longer need to write mlops for a new accelerator
@@ -137,9 +137,9 @@ Buffer # class of memory on this devic
unary_op (RELU, EXP, LOG, NEG, SIGN) # A -> A
reduce_op (SUM, MAX) # A -> B (smaller size, B has 1 in shape)
binary_op (ADD, SUB, MUL, DIV, POW, A, CMPEQ) # A + B -> C (broadcasting supported)
+reshape # A -> B (different size, same dim count)
perm_axis # A -> A
inner_slice # A -> B (different size)
-matmul # A + B -> C
conv, convdw, convdx # A + B -> C
```

MODIFY
tinygrad/llops/ops_cpu.py
tinygrad/llops/ops_cpu.py
@@ -63,12 +63,6 @@ def inner_slice(x, arg, ret):
ret[:] = x[tuple([slice(x[0], x[1], None) for x in slicee])]
return ret
-def matmul(a, b, c, transpose_a=False, transpose_b=False):
- if transpose_a: a = a.swapaxes(-2, -1)
- if transpose_b: b = b.swapaxes(-2, -1)
- c[:] = a @ b
- return c
-
def get_tx(x, C):
gx = x.reshape(C.bs,C.groups,C.cin,x.shape[2],x.shape[3])
return np.lib.stride_tricks.as_strided(gx,

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -192,42 +192,6 @@ def inner_slice(x, arg, ret):
buffer_np(np.array(shift, dtype=np.int32)))
return ret
-# c = a@b
-def matmul(a, b, c, transpose_a=False, transpose_b=False):
- cnt = np.prod(a.shape[0:-2]) if len(a.shape) > 2 else 1
- isize, msize, osize = i32(a.shape[-2]), i32(a.shape[-1]), i32(c.shape[-1])
- if transpose_a: isize,msize = msize,isize
- assert isize == c.shape[-2]
- assert (msize == b.shape[-1]) if transpose_b else (msize == b.shape[-2])
- assert (osize == b.shape[-2]) if transpose_b else (osize == b.shape[-1])
-
- matmul_prg = clbuild("matmul", """
- __kernel void matmul(
- __global const float *input, __global const float *weight, __global float *res,
- int isize, int is0, int is1, int msize, int ws0, int ws1, int osize
- ) {
- int stride = get_global_id(2);
-
- int X = get_global_id(0); // isize
- int Y = get_global_id(1); // osize
-
- float ret = 0.0;
- for (int x = 0; x < msize; x++) {
- ret += input[X * is0 + x * is1 + isize*msize*stride] * weight[Y * ws0 + x * ws1 + msize*osize*stride];
- }
-
- res[X * osize + Y + isize*osize*stride] = ret;
- }""")
-
- matmul_prg([isize, osize, cnt], None,
- a.cl, b.cl, c.cl,
- isize,
- msize if not transpose_a else i32(1), i32(1) if not transpose_a else isize,
- msize,
- i32(1) if not transpose_b else msize, osize if not transpose_b else i32(1),
- osize)
- return c
-
def conv(x,w,ret,stride,groups):
C = get_conv_args(x.shape, w.shape, stride, groups)
# input = (bs, groups, cin, iy, ix)

MODIFY
tinygrad/llops/ops_torch.py
tinygrad/llops/ops_torch.py
@@ -20,7 +20,7 @@ class TorchBuffer(torch.Tensor):
# ************* unary+binary+reduce+movement ops *************
-from tinygrad.llops.ops_cpu import unary_op, binary_op, reduce_op, reshape, perm_axis, inner_slice, matmul
+from tinygrad.llops.ops_cpu import unary_op, binary_op, reduce_op, reshape, perm_axis, inner_slice
# ************* processing ops *************

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -158,19 +158,6 @@ class Slice(Function):
# ************* processing ops *************
-class Matmul(Function):
- def forward(ctx, input, weight):
- assert input.shape[-1] == weight.shape[-2]
- ret = ctx.buffer(list(input.shape[0:-1])+[weight.shape[-1]])
- ctx.save_for_backward(input, weight)
- return ctx.op.matmul(input, weight, ret)
-
- def backward(ctx, grad_output):
- input, weight = ctx.saved_tensors
- grad_input = ctx.op.matmul(grad_output, weight, ctx.buffer(input.shape), transpose_b=True) if ctx.needs_input_grad[0] else None
- grad_weight = ctx.op.matmul(input, grad_output, ctx.buffer(weight.shape), transpose_a=True) if ctx.needs_input_grad[1] else None
- return grad_input, grad_weight
-
class Conv2D(Function):
def forward(ctx, x, w, stride=1, groups=1):
C = get_conv_args(x.shape, w.shape, stride, groups)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
import os, atexit, time, inspect, functools, importlib
from collections import defaultdict
import numpy as np
+from tinygrad.helpers import prod
# **** profiler ****
@@ -135,7 +136,7 @@ class Tensor:
@classmethod
def uniform(cls, *shape, **kwargs):
- return cls((np.random.uniform(-1., 1., size=shape)/np.sqrt(np.prod(shape))).astype(np.float32), **kwargs)
+ return cls((np.random.uniform(-1., 1., size=shape)/np.sqrt(prod(shape))).astype(np.float32), **kwargs)
@classmethod
def eye(cls, dim, **kwargs):
@@ -240,11 +241,20 @@ class Tensor:
def pad2d(self, padding):
return self[:, :, -padding[2]:self.shape[2]+padding[3], -padding[0]:self.shape[3]+padding[1]]
- def matmul(self, w):
- if len(self.shape) > 2 and len(w.shape) == 2:
- return self.reshape(shape=(-1, self.shape[-1]))._matmul(w).reshape(shape=list(self.shape[0:-1]) + [-1])
- else:
- return self._matmul(w)
+ def matmul(x, w):
+ bs, groups = prod(x.shape[0:-2]), prod(w.shape[0:-2])
+ cin, cout = w.shape[-2], w.shape[-1]
+ out_shape_t = tuple(list(x.shape[0:-2])+[cout,-1])
+ order = tuple(list(range(len(x.shape)-2))+[len(x.shape)-1, len(x.shape)-2])
+ worder = tuple(list(range(len(w.shape)-2))+[len(w.shape)-1, len(w.shape)-2])
+
+ # NOTE: with NHWC we can remove the transposes
+ # bs x groups*cin x H x W
+ cx = x.transpose(order=order).reshape(shape=(bs//groups, groups*cin, -1, 1))
+ # groups*cout x cin x H, W
+ cw = w.transpose(order=worder).reshape(shape=(groups*cout, cin, 1, 1))
+ return cx.conv2d(cw, groups=groups).reshape(shape=out_shape_t).transpose(order=order)
+
dot = matmul
def _canonicalize_reduce_axis(self, axis):
@@ -267,7 +277,7 @@ class Tensor:
def mean(self, axis=None, keepdim=False):
out = self.sum(axis=axis, keepdim=keepdim)
- return out * (np.prod(out.shape)/np.prod(self.shape))
+ return out * (prod(out.shape)/prod(self.shape))
def sqrt(self):
return self.pow(0.5)

Remove Matmul (#323)