factor out hand coded opt
MODIFY
extra/kernel_search.py
extra/kernel_search.py
@@ -45,6 +45,7 @@ def search(ast):
# get baseline
k = CLASTKernel(ast)
CL.time_sum = 0
+ #k.hand_coded_optimizations()
k.codegen()(*k.bufs)
winning_interventions = []
@@ -57,6 +58,7 @@ def search(ast):
inter = get_random_intervention(k)
apply_intervention(k, *inter)
+ k.simplify_ones()
# TODO: support upcasting, splitting, and local grouping for reduce
CL.time_sum = 0
@@ -83,7 +85,7 @@ def search(ast):
print(f"improved from {baseline/1e6:.2f} ms to {best_time/1e6:.2f} ms, a {baseline/best_time:.2f}x speedup")
if __name__ == "__main__":
- if int(os.getenv("OP", "0")):
+ if int(os.getenv("OP", "0")) == 1:
buf0 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 8, 4, 3, 3, 3, 4), views=[View((1, 130, 258, 1, 12), (393216, 3072, 12, 12, 1), -3084), ZeroView((1, 128, 256, 1, 12), ((0, 1), (-1, 129), (-1, 257), (0, 1), (0, 12))), View((1, 64, 128, 8, 4, 3, 3, 3, 4), (0, 6192, 24, 0, 0, 3096, 12, 4, 1), 0)]), hostbuf=GPUBuffer(shape=(128, 768, 4), force_create=True))
buf1 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 8, 4, 3, 3, 3, 4), views=[View((1, 64, 128, 8, 4, 3, 3, 3, 4), (0, 0, 0, 432, 4, 144, 16, 48, 1), 0)]), hostbuf=GPUBuffer(shape=(8, 108, 4), force_create=True))
op0 = LazyOp(BinaryOps.MUL, (buf0,buf1,), None)
@@ -99,6 +101,30 @@ if __name__ == "__main__":
op7 = LazyOp(BinaryOps.MUL, (buf3,op6,), None)
op8 = LazyOp(BinaryOps.SUB, (op3,op7,), None)
ast = LazyOp(MovementOps.RESHAPE, (op8,), (64, 1024, 4))
+ elif int(os.getenv("OP", "0")) == 2:
+ buf0 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 8, 4, 1, 1, 3, 3), views=[View((1, 66, 130, 32, 1), (262144, 4096, 32, 1, 1), -4128), ZeroView((1, 64, 128, 32, 1), ((0, 1), (-1, 65), (-1, 129), (0, 32), (0, 1))), View((1, 64, 128, 8, 4, 1, 1, 3, 3), (266240, 4160, 32, 4, 1, 12480, 12480, 4160, 32), 0)]), hostbuf=GPUBuffer(shape=(64, 1024, 4), force_create=True))
+ buf1 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 8, 4, 1, 1, 3, 3), views=[View((1, 64, 128, 8, 4, 1, 1, 3, 3), (0, 0, 0, 36, 1, 0, 0, 12, 4), 0)]), hostbuf=GPUBuffer(shape=(8, 9, 4), force_create=True))
+ op0 = LazyOp(BinaryOps.MUL, (buf0,buf1,), None)
+ op1 = LazyOp(ReduceOps.SUM, (op0,), (1, 64, 128, 8, 4, 1, 1, 1, 1))
+ buf2 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 8, 4, 1, 1, 1, 1), views=[View((1, 64, 128, 8, 4, 1, 1, 1, 1), (0, 0, 0, 4, 1, 1, 1, 1, 1), 0)]), hostbuf=GPUBuffer(shape=(32,), force_create=True))
+ op2 = LazyOp(BinaryOps.ADD, (op1,buf2,), None)
+ op3 = LazyOp(UnaryOps.RELU, (op2,), None)
+ buf3 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 8, 4, 1, 1, 1, 1), views=[View((1, 64, 128, 8, 4, 1, 1, 1, 1), (0, 0, 0, 0, 0, 1, 1, 1, 1), 0)]), hostbuf=GPUBuffer(shape=(1,), backing=np.array([1.], dtype=np.float32)))
+ buf4 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 8, 4, 1, 1, 1, 1), views=[View((1, 64, 128, 8, 4, 1, 1, 1, 1), (0, 0, 0, 0, 0, 1, 1, 1, 1), 0)]), hostbuf=GPUBuffer(shape=(1,), backing=np.array([1.], dtype=np.float32)))
+ op4 = LazyOp(UnaryOps.EXP, (op2,), None)
+ op5 = LazyOp(BinaryOps.SUB, (buf4,op4,), None)
+ op6 = LazyOp(UnaryOps.RELU, (op5,), None)
+ op7 = LazyOp(BinaryOps.MUL, (buf3,op6,), None)
+ op8 = LazyOp(BinaryOps.SUB, (op3,op7,), None)
+ ast = LazyOp(MovementOps.RESHAPE, (op8,), (64, 1024, 4))
+ elif int(os.getenv("OP", "0")) == 3:
+ buf0 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 4, 4, 1, 1, 8, 4), views=[View((1, 64, 128, 4, 4, 1, 1, 8, 4), (0, 4096, 32, 0, 0, 0, 0, 4, 1), 0)]), hostbuf=GPUBuffer(shape=(64, 1024, 4), force_create=True))
+ buf1 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 4, 4, 1, 1, 8, 4), views=[View((1, 64, 128, 4, 4, 1, 1, 8, 4), (0, 0, 0, 128, 4, 0, 0, 16, 1), 0)]), hostbuf=GPUBuffer(shape=(4, 32, 4), force_create=True))
+ op0 = LazyOp(BinaryOps.MUL, (buf0,buf1,), None)
+ op1 = LazyOp(ReduceOps.SUM, (op0,), (1, 64, 128, 4, 4, 1, 1, 1, 1))
+ buf2 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 4, 4, 1, 1, 1, 1), views=[View((1, 64, 128, 4, 4, 1, 1, 1, 1), (0, 0, 0, 4, 1, 1, 1, 1, 1), 0)]), hostbuf=GPUBuffer(shape=(16,), force_create=True))
+ op2 = LazyOp(BinaryOps.ADD, (op1,buf2,), None)
+ ast = LazyOp(MovementOps.RESHAPE, (op2,), (64, 512, 4))
elif int(os.getenv("BC", "0")):
# big conv
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))

MODIFY
tinygrad/ast.py
tinygrad/ast.py
@@ -59,6 +59,7 @@ class ASTKernel:
# process
# TODO: fetch from quick cache before processing
self.process()
+ self.group_for_reduce = []
def print(self):
buf_count = -1

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -206,15 +206,7 @@ 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):
- if DEBUG >= 3:
- print("old:", self.shapes)
- print("old:", self.strides)
-
- self.prekernel = set()
-
+ def hand_coded_optimizations(self):
# if there's images in the earlybufs, we have to make an axis the 4 loading one
# shove the axis to the end and remove
if any(isinstance(buf._buf, CLImage) for buf in self.earlybufs):
@@ -239,7 +231,6 @@ class CLASTKernel(ASTKernel):
self.simplify_ones()
# are we grouping?
- self.group_for_reduce = []
if self.buftokens[0].typ != Types.FLOAT4 and self.first_reduce <= 2 and self.first_reduce + 1 <= self.shape_len and prod(self.shapes[0][:self.first_reduce]) <= 2048:
for sz in ([256, 16] if prod(self.shapes[0][:self.first_reduce]) <= 32 else [16]):
if all([x[self.first_reduce] % sz == 0 or x[self.first_reduce] == 1 for x in self.shapes]):
@@ -309,7 +300,17 @@ class CLASTKernel(ASTKernel):
self.reshape_and_permute(lambda x: [base_shape[0], x[0]//base_shape[0]]+list(x[1:]), None)
self.simplify_ones()
+ # 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):
+ if DEBUG >= 3:
+ print("old:", self.shapes)
+ print("old:", self.strides)
+
+ self.hand_coded_optimizations()
+
# group for reduce
+ # TODO: clean this up to be in optimizations
self.output_shape = self.shapes[0][:self.first_reduce]
if len(self.group_for_reduce):
# with permute for memory coalesing
@@ -331,6 +332,7 @@ class CLASTKernel(ASTKernel):
self.bufs_to_delete : Set[int] = set()
self.loaded_keys : Dict[Tuple[int,int], Token] = {}
+ self.prekernel : Set[str] = set()
self.kernel : List[str] = ["const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n"]
# output_shape[-1] is get_global_id(0)
@@ -342,7 +344,7 @@ class CLASTKernel(ASTKernel):
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"]
- self.output_shape = [prod(self.output_shape[0:-2])] + self.output_shape[-2:]
+ self.output_shape = [prod(self.output_shape[0:-2])] + list(self.output_shape[-2:])
if DEBUG >= 3: print(f"replaced output shape with {self.output_shape}")
# early ast

factor out hand coded opt