first little attempt at search
ADD
extra/kernel_search.py
+#!/usr/bin/env python
+import random
+from tinygrad.ops import LazyOp, ReduceOps, BinaryOps, UnaryOps, MovementOps
+from tinygrad.shape import ShapeTracker, View, ZeroView
+from tinygrad.llops.ops_gpu import GPUBuffer, CLASTKernel
+
+def search(ast):
+ for i in range(20):
+ k = CLASTKernel(ast)
+
+ if i != 0:
+ # do search space
+ order_outer = list(range(0, k.first_reduce))
+ random.shuffle(order_outer)
+ order_inner = list(range(k.first_reduce, k.shape_len))
+ random.shuffle(order_inner)
+ order = order_outer + order_inner
+ print(order)
+ k.reshape_and_permute(None, order)
+
+ k.codegen()(*k.bufs)
+
+
+
+if __name__ == "__main__":
+ buf0 = GPUBuffer(shape=ShapeTracker(shape=(8, 1, 32, 112, 112, 3, 3, 3), views=[View((8, 3, 225, 225), (150528, 50176, 224, 1), 0), ZeroView((8, 3, 224, 224), ((0, 8), (0, 3), (0, 225), (0, 225))), View((8, 1, 32, 112, 112, 3, 3, 3), (151875, 151875, 0, 450, 2, 50625, 225, 1), 0)]), hostbuf=GPUBuffer(shape=(8, 3, 224, 224), force_create=True))
+ buf1 = GPUBuffer(shape=ShapeTracker(shape=(8, 1, 32, 112, 112, 3, 3, 3), views=[View((8, 1, 32, 112, 112, 3, 3, 3), (0, 0, 27, 0, 0, 9, 3, 1), 0)]), hostbuf=GPUBuffer(shape=(32, 3, 3, 3), force_create=True))
+ op0 = LazyOp(BinaryOps.MUL, (buf0,buf1,), None)
+ op1 = LazyOp(ReduceOps.SUM, (op0,), (8, 1, 32, 112, 112, 1, 1, 1))
+ ast = LazyOp(MovementOps.RESHAPE, (op1,), (8, 32, 112, 112))
+ search(ast)

MODIFY
test/external_test_gpu_ast.py
test/external_test_gpu_ast.py
@@ -100,6 +100,14 @@ class TestAST(unittest.TestCase):
ast = LazyOp(MovementOps.RESHAPE, (op1,), (1, 1, 1, 1, 4956))
compile_and_test_ast(ast)
+ def test_enet_first_conv_bs32(self):
+ buf0 = GPUBuffer(shape=ShapeTracker(shape=(8, 1, 32, 112, 112, 3, 3, 3), views=[View((8, 3, 225, 225), (150528, 50176, 224, 1), 0), ZeroView((8, 3, 224, 224), ((0, 8), (0, 3), (0, 225), (0, 225))), View((8, 1, 32, 112, 112, 3, 3, 3), (151875, 151875, 0, 450, 2, 50625, 225, 1), 0)]), hostbuf=GPUBuffer(shape=(8, 3, 224, 224), force_create=True))
+ buf1 = GPUBuffer(shape=ShapeTracker(shape=(8, 1, 32, 112, 112, 3, 3, 3), views=[View((8, 1, 32, 112, 112, 3, 3, 3), (0, 0, 27, 0, 0, 9, 3, 1), 0)]), hostbuf=GPUBuffer(shape=(32, 3, 3, 3), force_create=True))
+ op0 = LazyOp(BinaryOps.MUL, (buf0,buf1,), None)
+ op1 = LazyOp(ReduceOps.SUM, (op0,), (8, 1, 32, 112, 112, 1, 1, 1))
+ ast = LazyOp(MovementOps.RESHAPE, (op1,), (8, 32, 112, 112))
+ compile_and_test_ast(ast)
+
def test_enet_reduce_bs32(self):
buf0 = GPUBuffer(shape=ShapeTracker(shape=(3, 1, 32, 3, 3, 32, 112, 112), views=[View((3, 32, 225, 225), (50176, 150528, 224, 1), 0), ZeroView((3, 32, 224, 224), ((0, 3), (0, 32), (0, 225), (0, 225))), View((3, 1, 32, 3, 3, 32, 112, 112), (1620000, 1620000, 0, 225, 1, 50625, 450, 2), 0)]), hostbuf=GPUBuffer(shape=(32, 3, 224, 224), force_create=True))
buf1 = GPUBuffer(shape=ShapeTracker(shape=(3, 1, 32, 3, 3, 32, 112, 112), views=[View((3, 1, 32, 3, 3, 32, 112, 112), (0, 12845056, 401408, 0, 0, 12544, 112, 1), 0)]), hostbuf=GPUBuffer(shape=(1, 1, 32, 1, 1, 32, 112, 112), force_create=True))

MODIFY
tinygrad/ast.py
tinygrad/ast.py
@@ -56,6 +56,10 @@ class ASTKernel:
assert all_same([x.shape for x in self.bufs if x not in self.earlybufs]), "all latebufs must have the same shape"
assert all_same([len(x.shape) for x in self.bufs]), "all bufs must have the same shape size"
+ # process
+ # TODO: fetch from quick cache before processing
+ self.process()
+
def print(self):
buf_count = -1
op_count = -1

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -206,9 +206,9 @@ class CLASTKernel(ASTKernel):
else:
return [Token(code.replace("A", a.tok), a.typ) for a in values[0]]
+ # STOP WASTING TIME WITH DOING THE RESHAPES AND PERMUTES BY HAND. KERNEL SEARCH IS THE ONLY WAY IT WILL EVER BE GOOD
+ # group_for_reduce will have to be better first
def codegen(self):
- # TODO: fetch from quick cache before processing
- self.process()
if DEBUG >= 3:
print("old:", self.shapes)
print("old:", self.strides)
@@ -338,6 +338,7 @@ class CLASTKernel(ASTKernel):
if len(self.output_shape) > 3:
# sometimes, there's more dimensions. compact all the dimensions into the first one
+ # TODO: these compactions should be searchable
final_dimension = len(self.output_shape)-3
for i in range(len(self.output_shape)-4, -1, -1):
self.kernel += [f"int idx{i} = idx{final_dimension} % {self.output_shape[i]};", f"idx{final_dimension} = idx{final_dimension} / {self.output_shape[i]};\n"]

first little attempt at search