mawei/tinygrad
1
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-13 00:15:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

While fusion (#654)

Browse Source 
* try this

* readme

* opt comments
This commit is contained in:
George Hotz
2023-03-06 09:13:23 -08:00
committed by GitHub
parent 066a65dad5
commit 4b9bc1615b
2 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
README.md
View File

@@ -71,12 +71,14 @@ print(y.grad) # dz/dy
Try a matmul. See how, despite the style, it is fused into one kernel with the power of laziness.
```python
OPTLOCAL=1 GPU=1 DEBUG=3 python3 -c "from tinygrad.tensor import Tensor;
DEBUG=3 OPTLOCAL=1 GPU=1 python3 -c "from tinygrad.tensor import Tensor;
N = 1024; a, b = Tensor.randn(N, N), Tensor.randn(N, N);
c = (a.reshape(N, 1, N) * b.permute(1,0).reshape(1, N, N)).sum(axis=2);
print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
```
Change to `DEBUG=4` to see the generated code.
## Neural networks?
It turns out, a decent autograd tensor library is 90% of what you need for neural networks. Add an optimizer (SGD, RMSprop, and Adam implemented) from tinygrad.nn.optim, write some boilerplate minibatching code, and you have all you need.

8
tinygrad/codegen/gpu.py
View File

@@ -200,8 +200,8 @@ class GPUCodegen(ASTKernel):
# **** below this line need to be optional and benchmarked ****
# potentially do a second upcast based on a heuristic. this is optional and has nothing to do with float4
if prod(self.sts[0].shape[:self.first_reduce]) >= 1024:
# potentially do more upcasts of non reduce axes based on a heuristic
while prod(self.sts[0].shape[:self.first_reduce]) >= 1024:
xb_choices = []
for axis, upcast_amount in itertools.product(range(self.first_reduce), [3,4]): # consider all the non reduce axes, and a 3 or 4 reduce
# if it mods, and some buffer has stride 0 on axis while having no stride 0 in the buftoken
@@ -213,8 +213,10 @@ class GPUCodegen(ASTKernel):
self.shift_to(xb_choices[0][2], amount=xb_choices[0][3])
self.upcast()
self.simplify_ones()
else:
break
# if last dim <= 5 and it's a reduce dim, upcast (loop unrolling). no simplify needed since it's just an upcast. NOTE: careful, this has broken VALIDHACKS
# if last dim <= 5 and it's a reduce dim, upcast the reduce (loop unrolling). no simplify needed since it's just an upcast. NOTE: careful, this has broken VALIDHACKS
if self.first_reduce < self.shape_len and self.full_shape[-1] <= 5 and (max([x.size() for i,x in enumerate(self.buftokens) if self.bufs[i] in self.earlybufs]) <= 4 or not any(r for _,_,r in self.buftokens[self.full_buf_index].axis)):
self.upcast()