tests pass w/o float4

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)

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))

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] = {}