tests pass w/o float4
MODIFY
extra/kernel_search.py
extra/kernel_search.py
@@ -54,7 +54,7 @@ def apply_intervention(k, typ, dat):
lambda x: list(x[0:up_axis]) + ([x[up_axis]//amount, amount] if x[up_axis] > 1 else [1,1]) + list(x[up_axis+1:]),
[i for i in range(k.shape_len+1) if i != up_axis+1] + [up_axis+1])
# drop the last dimension
- k.upcast()
+ k.upcast(allow_float4=False)
elif typ == Interventions.SHIFT:
up_axis, amount, flip = dat[0], dat[1], dat[2]
k.reshape_and_permute(
@@ -139,6 +139,29 @@ def search(ast):
test_ast(k)
print(f"improved from {baseline/1e6:.2f} ms to {best_time/1e6:.2f} ms, a {baseline/best_time:.2f}x speedup @ {k.info.flops/best_time:.2f} GFLOPS")
+from tinygrad.ops import get_buffers
+def test_correctness(ast):
+ # before testing, we need to fill the buffers with randomness
+ bufs = get_buffers(ast)
+ for b in bufs:
+ randomness = np.random.default_rng().standard_normal(size=b._base_shape, dtype=np.float32)
+ if b._buf is not None: b._buf.copyin(randomness)
+ from extra.lib_test_ast import test_ast
+ k = CLASTKernel(ast)
+ ints = get_interventions(k)
+ k.codegen()(*k.bufs)
+ test_ast(k)
+ print("correct at baseline")
+ for int in ints:
+ print("***** APPLYING INTERVENTION", int)
+ k = CLASTKernel(ast)
+ k.printbufs("old:")
+ apply_intervention(k, *int)
+ k.printbufs("new:")
+ k.codegen()(*k.bufs)
+ print("***** TESTING INTERVENTION", int)
+ test_ast(k)
+
if __name__ == "__main__":
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))
@@ -212,4 +235,5 @@ if __name__ == "__main__":
op0 = LazyOp(BinaryOps.MUL, (buf0,buf1,), None)
op1 = LazyOp(ReduceOps.SUM, (op0,), (3, 1, 32, 3, 3, 1, 1, 1))
ast = LazyOp(MovementOps.RESHAPE, (op1,), (3, 32, 3, 3))
- search(ast)
+ #search(ast)
+ test_correctness(ast)

MODIFY
tinygrad/ast.py
tinygrad/ast.py
@@ -136,7 +136,7 @@ class ASTKernel:
if axis is not None: st.permute(*axis)
# drops the final dimension
- def upcast(self):
+ def upcast(self, allow_float4=True):
upcasted = [x.shape[-1] for x in self.sts if x.shape[-1] != 1]
assert len(upcasted) >= 1 and all_same(upcasted), f"can't upcast mismatch {upcasted}"
for i in range(len(self.bufs)):
@@ -144,7 +144,7 @@ class ASTKernel:
if st.shape[-1] == upcasted[0]:
# multiview shapetrackers can slice through a float4, so don't allow them
can_merge = (not st.needs_valid() and len(st.views) == 1) or "Image" in str(type(self.bufs[i]._buf)) # TODO: terrible hack
- if st.shape[-1] == 4 and self.buftokens[i].typ == Types.FLOAT and st.views[-1].strides[-1] == 1 and can_merge:
+ if allow_float4 and st.shape[-1] == 4 and self.buftokens[i].typ == Types.FLOAT and st.views[-1].strides[-1] == 1 and can_merge:
# this is an upcast to FLOAT4
self.buftokens[i].typ = Types.FLOAT4
assert all(st.views[-1].strides[i]%upcasted[0] == 0 or st.views[-1].shape[i] == 1 for i in range(len(st.shape)-1))

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -220,13 +220,14 @@ class CLASTKernel(ASTKernel):
if self.first_reduce < self.shape_len and end_dimension > 1 and end_dimension <= 3 and max([x.size() for i,x in enumerate(self.buftokens) if self.bufs[i] in self.earlybufs]) <= 4:
self.upcast()
+ def printbufs(self, prefix=""):
+ print(f"first_reduce: {self.first_reduce} shape_len: {self.shape_len} group_for_reduce: {self.group_for_reduce}")
+ for i in range(len(self.sts)):
+ print(prefix, self.buftokens[i], f"early:{'T' if i < len(self.bufs) and self.bufs[i] in self.earlybufs else 'F'} image:{'T' if i < len(self.bufs) and isinstance(self.bufs[i]._buf, CLImage) else 'F'}", self.sts[i].shape, self.sts[i].views[-1].strides)
+
# 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:", [x.shape for x in self.sts])
- print("old:", [x.views[-1].strides for x in self.sts])
-
self.hand_coded_optimizations()
# add a local buffer for multistage reduce
@@ -236,10 +237,8 @@ class CLASTKernel(ASTKernel):
self.output_shape = list(self.sts[0].shape[:self.first_reduce]) + self.group_for_reduce
if DEBUG >= 3:
- print(f"first_reduce: {self.first_reduce} shape_len: {self.shape_len} group_for_reduce: {self.group_for_reduce}")
print("output shape", self.output_shape)
- for i in range(len(self.sts)):
- print(self.buftokens[i], f"early:{'T' if i < len(self.bufs) and self.bufs[i] in self.earlybufs else 'F'} image:{'T' if i < len(self.bufs) and isinstance(self.bufs[i]._buf, CLImage) else 'F'}", self.sts[i])
+ self.printbufs("new:")
self.bufs_to_delete : Set[int] = set()
self.loaded_keys : Dict[Tuple[int,int], Token] = {}

tests pass w/o float4