Merge branch 'master' into sched_lin

* fix cast folding in llama

* dtypes that work everywhere

* Skip test_cast_there_and_back for backend casts

Skip test due to backend casting issues.

test/test_linearizer.py

@@ -14,6 +14,8 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
 MOCKGPU = getenv("MOCKGPU")
 
+from tinygrad.uop.ops import print_uops # noqa: F401 # pylint: disable=unused-import
+
 class TestLinearizer(unittest.TestCase):
   def test_arg_dedup(self):
     # NOTE: this realize exists because Tensor.numpy calls .contiguous() internally
@@ -38,6 +40,22 @@ class TestLinearizer(unittest.TestCase):
     np.testing.assert_equal(a.numpy(), ta)
     np.testing.assert_equal(b.numpy(), tb)
 
+  @unittest.skip("TODO: some backends insert more casts")
+  def test_cast_there_and_back(self):
+    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
+    out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
+    ast = helper_linearizer_opt(out)
+    uops = get_program(ast, opts=[]).uops
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
+
+  @unittest.expectedFailure
+  def test_cast_back_and_there(self):
+    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
+    out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
+    ast = helper_linearizer_opt(out)
+    uops = get_program(ast, opts=[]).uops
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
+
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
   def test_late_bias_load(self):
     img = Tensor.empty(1, 3, 16, 16)
@@ -491,7 +509,7 @@ def copyout_outputs(outbufs:list[Buffer]) -> list[np.ndarray]:
   return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
 
 def reset_bufs(bufs:list[Buffer]):
-  for buf in bufs: buf.copyin(np.zeros((buf.size, ), dtype=_to_np_dtype(buf.dtype)).data) # Zero to check that all values are filled
+  for buf in bufs: buf.copyin(np.zeros((buf.size*buf.dtype.itemsize,), dtype=np.uint8).data)
 
 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
                                apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):

tinygrad/codegen/late/linearizer.py

@@ -4,7 +4,7 @@ from collections import defaultdict
 from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
 from tinygrad.helpers import prod, getenv, TUPLE_ORDER
 
-def linearize(sink:UOp) -> list[UOp]:
+def linearize(sink:UOp, tuple_order=TUPLE_ORDER) -> list[UOp]:
   # this is a toposort with priority
   lst = list(sink.toposort())
   consumers: defaultdict[UOp, list[UOp]] = defaultdict(list)
@@ -39,7 +39,7 @@ def linearize(sink:UOp) -> list[UOp]:
     priorities[u] = (run_count, priority, extra)
 
   # number the uops in "ideal" order
-  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}
+  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if tuple_order else ())))}
 
   # then force them to be toposorted in as close to the ideal order as possible
   heap = [(-nkey[sink], sink)]

tinygrad/engine/schedule.py

@@ -1,9 +1,10 @@
 from typing import cast
 from dataclasses import dataclass, field
 from collections import deque, defaultdict
-from tinygrad.uop.ops import UOp, Ops, buffers
+from tinygrad.uop.ops import UOp, Ops, buffers, print_uops
 from tinygrad.device import Device, Buffer, MultiBuffer
 from tinygrad.helpers import Metadata, all_same
+from tinygrad.codegen.late.linearizer import linearize
 
 # **** ScheduleItem return type
 
@@ -17,6 +18,19 @@ class ScheduleItem:
 # **** schedule linearizer
 
 def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[str, int]]:
+  lst = linearize(sched_sink, tuple_order=False)
+  print_uops(lst)
+
+  schedule: list[ScheduleItem] = []
+  var_vals: dict[str, int] = {}
+  for k in lst:
+    if k.op is Ops.KERNEL:
+      ubufs = tuple(s.buf_uop.buffer for s in k.src if s.op is not Ops.BIND)
+      # ONE -> ONE
+      schedule.append(ScheduleItem(k.arg.ast, cast(tuple[Buffer, ...], ubufs), k.arg.metadata))
+      pass
+
+  """
   # construct the KERNEL children graph based on assigns
   children: defaultdict[UOp, list[UOp]] = defaultdict(list)
   in_degree: dict[UOp, int] = {}
@@ -79,5 +93,6 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
     for x in children[k]:
       in_degree[x] -= 1
       if in_degree[x] == 0: queues[_heuristic(x)].append(x)
+  """
 
   return schedule, var_vals

tinygrad/uop/ops.py

@@ -860,7 +860,8 @@ def print_uops(uops:list[UOp]):
   uops_index = {u:i for i,u in enumerate(uops)}
   for i,u in enumerate(uops):
     formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
-    print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
+    formatted_arg = str(u.arg)[0:30].replace("\n", "")
+    print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {formatted_arg}")
 
 # ***** pattern matcher *****