We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
LLVM test failing without OPT
OPT=1 LLVM=1 python3 test/test_ops.py TestOps.test_topo_sort

Still failing, but default OPT is now 2 so you have to say OPT=1 to make it fail.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tinygrad/ast.py`
Content:
```
1 from enum import Enum, auto
2 import itertools
3 from typing import List, Tuple, Optional
4 from tinygrad.helpers import prod, dedup, all_same
5 from tinygrad.ops import LazyOp, MovementOps, get_lazyop_info, get_buffers, ReduceOps, get_lazyops
6 from tinygrad.shape import ShapeTracker, View, strides_for_shape
7 
8 def get_first_reduce(shapes):
9   for i in range(len(shapes[0])):
10     if not all_same([x[i] for x in shapes]):
11       return i
12   return len(shapes[0])  # off the end
13 
14 # this will be removed soon anyway
15 class Types(Enum): FLOAT = auto(); FLOAT4 = auto() # noqa: E702
16 class Token:
17   def __init__(self, tok:str, typ:Types, ptr:bool=False):
18     assert isinstance(tok, str)
19     self.tok, self.typ, self.ptr = tok, typ, ptr
20     self.axis : List[Tuple[int, int, bool]] = []
21   def array(self, length, stride, reduce): self.axis.append((length, stride, reduce))
22   def size(self): return prod([x[0] for x in self.axis])
23   def offsets(self): return [sum(t) for t in itertools.product(*[[y*x[1] for y in range(x[0])] for x in self.axis[::-1]])] if len(self.axis) else [0]
24   # TODO: this is sort of a hack, it gets the accumulator indices
25   def acc_offsets(self):
26     if len(self.axis) == 0: return [0]
27     acc_strides = [x*(1-self.axis[::-1][i][2]) for i,x in enumerate(strides_for_shape(tuple(1 if r else s for s,_,r in self.axis[::-1])))]
28     return [sum(t) for t in itertools.product(*[[y*acc_strides[i] for y in range(x[0])] for i,x in enumerate(self.axis[::-1])])]
29   def decltype(self): return ('float' if self.typ == Types.FLOAT else 'float4') + ('*' if self.ptr else str())
30   def __repr__(self): return f"<{self.typ}{'*' if self.ptr else str()} {self.tok}{f'[{self.axis}]' if len(self.axis) else str()}>"
31 
32 # ast kernel can contain one ReduceOp with arbitrary Binary/Unary ops
33 class ASTKernel:
34   def __init__(self, ast:LazyOp, output_buffer=None):
35     # key for lookup in cache (can change, str might not be right)
36     self.input_ast = ast
37     self.key = str(ast)
38 
39     # if the AST ends with a RESHAPE, we remove it and create the buffer accordingly
40     if ast.op == MovementOps.RESHAPE:
41       output_shape = ast.arg
42       ast = ast.src[0]
43     else:
44       output_shape = None
45 
46     self.info = get_lazyop_info(ast)
47     self.bufs = dedup(get_buffers(ast))
48     self.ast = ast
49 
50     # check if the output buffer is allowed to be used
51     # if it's aliased, don't use it
52     if output_buffer is not None:
53       for a in self.bufs:
54         if a._buf == output_buffer._buf and not a.st.contiguous:
55           output_buffer = None
56           break
57 
58     # create the buffer we are returning (as the same type as the input buffers) and add it as the first buffer
59     self.ret = output_buffer if output_buffer else type(self.bufs[0])(output_shape if output_shape else self.info.shape, force_create=True)
60     self.bufs = ([type(self.ret)(self.info.shape, hostbuf=self.ret)] if output_shape else [self.ret]) + self.bufs
61 
62   def process(self) -> None:
63     if hasattr(self, "sts"): return   # already processed
64 
65     reduceops = [x for x in get_lazyops(self.ast) if x.op in ReduceOps]
66     assert len(dedup(reduceops)) <= 1, "max one reduce op in an ast"
67     self.reduceop = reduceops[0] if reduceops else None
68     self.reduceopop : Optional[ReduceOps] = self.reduceop.op if self.reduceop is not None and isinstance(self.reduceop.op, ReduceOps) else None
69     self.earlybufs = dedup(get_buffers(self.reduceop)) if self.reduceop else []
70 
71     self.buftokens = [Token(f"data{i}", Types.FLOAT, ptr=True) for i in range(len(self.bufs))]
72     self.group_for_reduce : List[int] = []
73 
74     # check valid AST kernel
75     assert all_same([x.shape for x in self.earlybufs]), "all earlybufs must have the same shape"
76     assert all_same([x.shape for x in self.bufs if x not in self.earlybufs]), "all latebufs must have the same shape"
77     assert all_same([len(x.shape) for x in self.bufs]), "all bufs must have the same shape size"
78 
79     # process
80     self.sts : List[ShapeTracker] = [x.st.copy() for x in self.bufs]   # create new shapetrackers inside this kernel
81     self.simplify_ones()
82     self.simplify_merge_adjacent()
83 
84   def print(self):
85     buf_count = -1
86     op_count = -1
87     cache = {}
88     def print_ast(x, name=None):
89       nonlocal buf_count, op_count
90       if x not in cache:
91         if not isinstance(x, LazyOp):
92           if name is None:
93             buf_count += 1
94             name = f"buf{buf_count}"
95           print(f"buf{buf_count} = {x}")
96           cache[x] = name
97         else:
98           srcs = [print_ast(y) for y in x.src]
99           if name is None:
100             op_count += 1
101             name = f"op{op_count}"
102           print(f"{name} = LazyOp({str(x.op)}, ({','.join(srcs)},), {x.arg})")
103           cache[x] = name
104       return cache[x]
105     print_ast(self.input_ast, "ast")
106 
107   def printbufs(self, prefix=""):
108     print(f"first_reduce: {self.first_reduce} shape_len: {self.shape_len} group_for_reduce: {self.group_for_reduce}")
109     for i in range(len(self.sts)):
110       print(prefix, self.buftokens[i], f"early:{'T' if i < len(self.bufs) and self.bufs[i] in self.earlybufs else 'F'}", self.sts[i].shape, self.sts[i].views[-1].strides, len(self.sts[i].views), type(self.bufs[i]._buf))
111 
112   @property
113   def shape_len(self) -> int: return len(self.sts[0].shape)
114 
115   def simplify_ones(self):
116     # remove places where the shape is all ones
117     # TODO: this should be factored in to multi shape stride
118     all_ones = [all(st.shape[i]==1 for st in self.sts) for i in range(self.shape_len)]
119     # keep at least 1 one
120     if all(all_ones): all_ones[-1] = False
121     self.reshape_and_permute(lambda shape: [x for i,x in enumerate(shape) if not all_ones[i]], None)
122     # find first mismatch, don't reduce this
123     self.first_reduce = get_first_reduce([x.shape for x in self.sts])
124 
125   def simplify_merge_adjacent(self):
126     shapes, strides = [x.shape for x in self.sts], [x.views[-1].strides for x in self.sts]
127 
128     # merge dimensions if we can, multi get_shape_strides
129     # TODO: does this always preserve the reduce dimension, NO
130     # TODO: move this into shapetracker, with tests!
131     rets = [[(shapes[j][0], strides[j][0])] for j in range(len(shapes))]
132     for i in range(1, len(shapes[0])):
133       can_merge = []
134       for j in range(len(shapes)):
135         # TODO: added the always mergeability of 1s, is this right? if so, add to shapetracker in the 1 case
136         can_merge.append((strides[j][i] != 0 and rets[j][-1][1] == shapes[j][i]*strides[j][i]) or (strides[j][i] == 0 and rets[j][-1][1] == 0))
137       # more can merge than this
138       mergeable = all(can_merge) and i != self.first_reduce
139       for j in range(len(shapes)):
140         if mergeable:
141           rets[j][-1] = (rets[j][-1][0] * shapes[j][i], strides[j][i])
142         else:
143           rets[j].append((shapes[j][i], strides[j][i]))
144 
145     for i,x in enumerate(rets): self.sts[i].reshape(tuple(y[0] for y in x))
146     self.first_reduce = get_first_reduce([x.shape for x in self.sts])
147 
148   # this should be aware of the three parts to the shape
149   #  * the input/output dimensions
150   #  * the reduce dimensions
151   #  * the size outputted by each kernel
152   def reshape_and_permute(self, new_shape_fxn, axis):
153     for st in self.sts:
154       if new_shape_fxn is not None: st.reshape(tuple(new_shape_fxn(st.shape)))
155       if axis is not None: st.permute(tuple(axis))
156 
157   # drops the final dimension
158   def upcast(self, allow_float4=True):
159     upcasted = [x.shape[-1] for x in self.sts if x.shape[-1] != 1]
160     assert len(upcasted) >= 1 and all_same(upcasted), f"can't upcast mismatch {upcasted}"
161     for i in range(len(self.bufs)):
162       st = self.sts[i]
163       if st.shape[-1] == upcasted[0]:
164         # multiview shapetrackers can slice through a float4, so don't allow them
165         can_merge = (not st.needs_valid() and len(st.views) == 1) or "Image" in str(type(self.bufs[i]._buf))  # TODO: terrible hack
166         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:
167           # this is an upcast to FLOAT4
168           self.buftokens[i].typ = Types.FLOAT4
169           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))
170           assert self.sts[i].offset % upcasted[0] == 0
171         else:
172           self.buftokens[i].array(upcasted[0], st.views[-1].strides[-1], len(upcasted) != len(self.sts))
173 
174     # remove the last dimension
175     for st in self.sts: st.views[-1] = View(st.shape[0:-1], st.views[-1].strides[0:-1], st.views[-1].offset)
176
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/tinygrad/ast.py b/tinygrad/ast.py
--- a/tinygrad/ast.py
+++ b/tinygrad/ast.py
@@ -32,9 +32,7 @@
 # ast kernel can contain one ReduceOp with arbitrary Binary/Unary ops
 class ASTKernel:
   def __init__(self, ast:LazyOp, output_buffer=None):
-    # key for lookup in cache (can change, str might not be right)
     self.input_ast = ast
-    self.key = str(ast)
 
     # if the AST ends with a RESHAPE, we remove it and create the buffer accordingly
     if ast.op == MovementOps.RESHAPE:
@@ -59,6 +57,10 @@
     self.ret = output_buffer if output_buffer else type(self.bufs[0])(output_shape if output_shape else self.info.shape, force_create=True)
     self.bufs = ([type(self.ret)(self.info.shape, hostbuf=self.ret)] if output_shape else [self.ret]) + self.bufs
 
+    # key for lookup in cache (can change, str might not be right)
+    # bufs are needed because kernels like f(x) = x + x and f(x, y) = x + y have the same str(ast), but are different kernels.
+    self.key = f"ASTKernelKey ast={str(ast)} bufs={self.bufs}"
+
   def process(self) -> None:
     if hasattr(self, "sts"): return   # already processed

{"golden_diff": "diff --git a/tinygrad/ast.py b/tinygrad/ast.py\n--- a/tinygrad/ast.py\n+++ b/tinygrad/ast.py\n@@ -32,9 +32,7 @@\n # ast kernel can contain one ReduceOp with arbitrary Binary/Unary ops\n class ASTKernel:\n   def __init__(self, ast:LazyOp, output_buffer=None):\n-    # key for lookup in cache (can change, str might not be right)\n     self.input_ast = ast\n-    self.key = str(ast)\n \n     # if the AST ends with a RESHAPE, we remove it and create the buffer accordingly\n     if ast.op == MovementOps.RESHAPE:\n@@ -59,6 +57,10 @@\n     self.ret = output_buffer if output_buffer else type(self.bufs[0])(output_shape if output_shape else self.info.shape, force_create=True)\n     self.bufs = ([type(self.ret)(self.info.shape, hostbuf=self.ret)] if output_shape else [self.ret]) + self.bufs\n \n+    # key for lookup in cache (can change, str might not be right)\n+    # bufs are needed because kernels like f(x) = x + x and f(x, y) = x + y have the same str(ast), but are different kernels.\n+    self.key = f\"ASTKernelKey ast={str(ast)} bufs={self.bufs}\"\n+\n   def process(self) -> None:\n     if hasattr(self, \"sts\"): return   # already processed\n", "issue": "LLVM test failing without OPT\nOPT=1 LLVM=1 python3 test/test_ops.py TestOps.test_topo_sort\r\n\r\nStill failing, but default OPT is now 2 so you have to say OPT=1 to make it fail.\n", "before_files": [{"content": "from enum import Enum, auto\nimport itertools\nfrom typing import List, Tuple, Optional\nfrom tinygrad.helpers import prod, dedup, all_same\nfrom tinygrad.ops import LazyOp, MovementOps, get_lazyop_info, get_buffers, ReduceOps, get_lazyops\nfrom tinygrad.shape import ShapeTracker, View, strides_for_shape\n\ndef get_first_reduce(shapes):\n  for i in range(len(shapes[0])):\n    if not all_same([x[i] for x in shapes]):\n      return i\n  return len(shapes[0])  # off the end\n\n# this will be removed soon anyway\nclass Types(Enum): FLOAT = auto(); FLOAT4 = auto() # noqa: E702\nclass Token:\n  def __init__(self, tok:str, typ:Types, ptr:bool=False):\n    assert isinstance(tok, str)\n    self.tok, self.typ, self.ptr = tok, typ, ptr\n    self.axis : List[Tuple[int, int, bool]] = []\n  def array(self, length, stride, reduce): self.axis.append((length, stride, reduce))\n  def size(self): return prod([x[0] for x in self.axis])\n  def offsets(self): return [sum(t) for t in itertools.product(*[[y*x[1] for y in range(x[0])] for x in self.axis[::-1]])] if len(self.axis) else [0]\n  # TODO: this is sort of a hack, it gets the accumulator indices\n  def acc_offsets(self):\n    if len(self.axis) == 0: return [0]\n    acc_strides = [x*(1-self.axis[::-1][i][2]) for i,x in enumerate(strides_for_shape(tuple(1 if r else s for s,_,r in self.axis[::-1])))]\n    return [sum(t) for t in itertools.product(*[[y*acc_strides[i] for y in range(x[0])] for i,x in enumerate(self.axis[::-1])])]\n  def decltype(self): return ('float' if self.typ == Types.FLOAT else 'float4') + ('*' if self.ptr else str())\n  def __repr__(self): return f\"<{self.typ}{'*' if self.ptr else str()} {self.tok}{f'[{self.axis}]' if len(self.axis) else str()}>\"\n\n# ast kernel can contain one ReduceOp with arbitrary Binary/Unary ops\nclass ASTKernel:\n  def __init__(self, ast:LazyOp, output_buffer=None):\n    # key for lookup in cache (can change, str might not be right)\n    self.input_ast = ast\n    self.key = str(ast)\n\n    # if the AST ends with a RESHAPE, we remove it and create the buffer accordingly\n    if ast.op == MovementOps.RESHAPE:\n      output_shape = ast.arg\n      ast = ast.src[0]\n    else:\n      output_shape = None\n\n    self.info = get_lazyop_info(ast)\n    self.bufs = dedup(get_buffers(ast))\n    self.ast = ast\n\n    # check if the output buffer is allowed to be used\n    # if it's aliased, don't use it\n    if output_buffer is not None:\n      for a in self.bufs:\n        if a._buf == output_buffer._buf and not a.st.contiguous:\n          output_buffer = None\n          break\n\n    # create the buffer we are returning (as the same type as the input buffers) and add it as the first buffer\n    self.ret = output_buffer if output_buffer else type(self.bufs[0])(output_shape if output_shape else self.info.shape, force_create=True)\n    self.bufs = ([type(self.ret)(self.info.shape, hostbuf=self.ret)] if output_shape else [self.ret]) + self.bufs\n\n  def process(self) -> None:\n    if hasattr(self, \"sts\"): return   # already processed\n\n    reduceops = [x for x in get_lazyops(self.ast) if x.op in ReduceOps]\n    assert len(dedup(reduceops)) <= 1, \"max one reduce op in an ast\"\n    self.reduceop = reduceops[0] if reduceops else None\n    self.reduceopop : Optional[ReduceOps] = self.reduceop.op if self.reduceop is not None and isinstance(self.reduceop.op, ReduceOps) else None\n    self.earlybufs = dedup(get_buffers(self.reduceop)) if self.reduceop else []\n\n    self.buftokens = [Token(f\"data{i}\", Types.FLOAT, ptr=True) for i in range(len(self.bufs))]\n    self.group_for_reduce : List[int] = []\n\n    # check valid AST kernel\n    assert all_same([x.shape for x in self.earlybufs]), \"all earlybufs must have the same shape\"\n    assert all_same([x.shape for x in self.bufs if x not in self.earlybufs]), \"all latebufs must have the same shape\"\n    assert all_same([len(x.shape) for x in self.bufs]), \"all bufs must have the same shape size\"\n\n    # process\n    self.sts : List[ShapeTracker] = [x.st.copy() for x in self.bufs]   # create new shapetrackers inside this kernel\n    self.simplify_ones()\n    self.simplify_merge_adjacent()\n\n  def print(self):\n    buf_count = -1\n    op_count = -1\n    cache = {}\n    def print_ast(x, name=None):\n      nonlocal buf_count, op_count\n      if x not in cache:\n        if not isinstance(x, LazyOp):\n          if name is None:\n            buf_count += 1\n            name = f\"buf{buf_count}\"\n          print(f\"buf{buf_count} = {x}\")\n          cache[x] = name\n        else:\n          srcs = [print_ast(y) for y in x.src]\n          if name is None:\n            op_count += 1\n            name = f\"op{op_count}\"\n          print(f\"{name} = LazyOp({str(x.op)}, ({','.join(srcs)},), {x.arg})\")\n          cache[x] = name\n      return cache[x]\n    print_ast(self.input_ast, \"ast\")\n\n  def printbufs(self, prefix=\"\"):\n    print(f\"first_reduce: {self.first_reduce} shape_len: {self.shape_len} group_for_reduce: {self.group_for_reduce}\")\n    for i in range(len(self.sts)):\n      print(prefix, self.buftokens[i], f\"early:{'T' if i < len(self.bufs) and self.bufs[i] in self.earlybufs else 'F'}\", self.sts[i].shape, self.sts[i].views[-1].strides, len(self.sts[i].views), type(self.bufs[i]._buf))\n\n  @property\n  def shape_len(self) -> int: return len(self.sts[0].shape)\n\n  def simplify_ones(self):\n    # remove places where the shape is all ones\n    # TODO: this should be factored in to multi shape stride\n    all_ones = [all(st.shape[i]==1 for st in self.sts) for i in range(self.shape_len)]\n    # keep at least 1 one\n    if all(all_ones): all_ones[-1] = False\n    self.reshape_and_permute(lambda shape: [x for i,x in enumerate(shape) if not all_ones[i]], None)\n    # find first mismatch, don't reduce this\n    self.first_reduce = get_first_reduce([x.shape for x in self.sts])\n\n  def simplify_merge_adjacent(self):\n    shapes, strides = [x.shape for x in self.sts], [x.views[-1].strides for x in self.sts]\n\n    # merge dimensions if we can, multi get_shape_strides\n    # TODO: does this always preserve the reduce dimension, NO\n    # TODO: move this into shapetracker, with tests!\n    rets = [[(shapes[j][0], strides[j][0])] for j in range(len(shapes))]\n    for i in range(1, len(shapes[0])):\n      can_merge = []\n      for j in range(len(shapes)):\n        # TODO: added the always mergeability of 1s, is this right? if so, add to shapetracker in the 1 case\n        can_merge.append((strides[j][i] != 0 and rets[j][-1][1] == shapes[j][i]*strides[j][i]) or (strides[j][i] == 0 and rets[j][-1][1] == 0))\n      # more can merge than this\n      mergeable = all(can_merge) and i != self.first_reduce\n      for j in range(len(shapes)):\n        if mergeable:\n          rets[j][-1] = (rets[j][-1][0] * shapes[j][i], strides[j][i])\n        else:\n          rets[j].append((shapes[j][i], strides[j][i]))\n\n    for i,x in enumerate(rets): self.sts[i].reshape(tuple(y[0] for y in x))\n    self.first_reduce = get_first_reduce([x.shape for x in self.sts])\n\n  # this should be aware of the three parts to the shape\n  #  * the input/output dimensions\n  #  * the reduce dimensions\n  #  * the size outputted by each kernel\n  def reshape_and_permute(self, new_shape_fxn, axis):\n    for st in self.sts:\n      if new_shape_fxn is not None: st.reshape(tuple(new_shape_fxn(st.shape)))\n      if axis is not None: st.permute(tuple(axis))\n\n  # drops the final dimension\n  def upcast(self, allow_float4=True):\n    upcasted = [x.shape[-1] for x in self.sts if x.shape[-1] != 1]\n    assert len(upcasted) >= 1 and all_same(upcasted), f\"can't upcast mismatch {upcasted}\"\n    for i in range(len(self.bufs)):\n      st = self.sts[i]\n      if st.shape[-1] == upcasted[0]:\n        # multiview shapetrackers can slice through a float4, so don't allow them\n        can_merge = (not st.needs_valid() and len(st.views) == 1) or \"Image\" in str(type(self.bufs[i]._buf))  # TODO: terrible hack\n        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:\n          # this is an upcast to FLOAT4\n          self.buftokens[i].typ = Types.FLOAT4\n          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))\n          assert self.sts[i].offset % upcasted[0] == 0\n        else:\n          self.buftokens[i].array(upcasted[0], st.views[-1].strides[-1], len(upcasted) != len(self.sts))\n\n    # remove the last dimension\n    for st in self.sts: st.views[-1] = View(st.shape[0:-1], st.views[-1].strides[0:-1], st.views[-1].offset)\n", "path": "tinygrad/ast.py"}], "after_files": [{"content": "from enum import Enum, auto\nimport itertools\nfrom typing import List, Tuple, Optional\nfrom tinygrad.helpers import prod, dedup, all_same\nfrom tinygrad.ops import LazyOp, MovementOps, get_lazyop_info, get_buffers, ReduceOps, get_lazyops\nfrom tinygrad.shape import ShapeTracker, View, strides_for_shape\n\ndef get_first_reduce(shapes):\n  for i in range(len(shapes[0])):\n    if not all_same([x[i] for x in shapes]):\n      return i\n  return len(shapes[0])  # off the end\n\n# this will be removed soon anyway\nclass Types(Enum): FLOAT = auto(); FLOAT4 = auto() # noqa: E702\nclass Token:\n  def __init__(self, tok:str, typ:Types, ptr:bool=False):\n    assert isinstance(tok, str)\n    self.tok, self.typ, self.ptr = tok, typ, ptr\n    self.axis : List[Tuple[int, int, bool]] = []\n  def array(self, length, stride, reduce): self.axis.append((length, stride, reduce))\n  def size(self): return prod([x[0] for x in self.axis])\n  def offsets(self): return [sum(t) for t in itertools.product(*[[y*x[1] for y in range(x[0])] for x in self.axis[::-1]])] if len(self.axis) else [0]\n  # TODO: this is sort of a hack, it gets the accumulator indices\n  def acc_offsets(self):\n    if len(self.axis) == 0: return [0]\n    acc_strides = [x*(1-self.axis[::-1][i][2]) for i,x in enumerate(strides_for_shape(tuple(1 if r else s for s,_,r in self.axis[::-1])))]\n    return [sum(t) for t in itertools.product(*[[y*acc_strides[i] for y in range(x[0])] for i,x in enumerate(self.axis[::-1])])]\n  def decltype(self): return ('float' if self.typ == Types.FLOAT else 'float4') + ('*' if self.ptr else str())\n  def __repr__(self): return f\"<{self.typ}{'*' if self.ptr else str()} {self.tok}{f'[{self.axis}]' if len(self.axis) else str()}>\"\n\n# ast kernel can contain one ReduceOp with arbitrary Binary/Unary ops\nclass ASTKernel:\n  def __init__(self, ast:LazyOp, output_buffer=None):\n    self.input_ast = ast\n\n    # if the AST ends with a RESHAPE, we remove it and create the buffer accordingly\n    if ast.op == MovementOps.RESHAPE:\n      output_shape = ast.arg\n      ast = ast.src[0]\n    else:\n      output_shape = None\n\n    self.info = get_lazyop_info(ast)\n    self.bufs = dedup(get_buffers(ast))\n    self.ast = ast\n\n    # check if the output buffer is allowed to be used\n    # if it's aliased, don't use it\n    if output_buffer is not None:\n      for a in self.bufs:\n        if a._buf == output_buffer._buf and not a.st.contiguous:\n          output_buffer = None\n          break\n\n    # create the buffer we are returning (as the same type as the input buffers) and add it as the first buffer\n    self.ret = output_buffer if output_buffer else type(self.bufs[0])(output_shape if output_shape else self.info.shape, force_create=True)\n    self.bufs = ([type(self.ret)(self.info.shape, hostbuf=self.ret)] if output_shape else [self.ret]) + self.bufs\n\n    # key for lookup in cache (can change, str might not be right)\n    # bufs are needed because kernels like f(x) = x + x and f(x, y) = x + y have the same str(ast), but are different kernels.\n    self.key = f\"ASTKernelKey ast={str(ast)} bufs={self.bufs}\"\n\n  def process(self) -> None:\n    if hasattr(self, \"sts\"): return   # already processed\n\n    reduceops = [x for x in get_lazyops(self.ast) if x.op in ReduceOps]\n    assert len(dedup(reduceops)) <= 1, \"max one reduce op in an ast\"\n    self.reduceop = reduceops[0] if reduceops else None\n    self.reduceopop : Optional[ReduceOps] = self.reduceop.op if self.reduceop is not None and isinstance(self.reduceop.op, ReduceOps) else None\n    self.earlybufs = dedup(get_buffers(self.reduceop)) if self.reduceop else []\n\n    self.buftokens = [Token(f\"data{i}\", Types.FLOAT, ptr=True) for i in range(len(self.bufs))]\n    self.group_for_reduce : List[int] = []\n\n    # check valid AST kernel\n    assert all_same([x.shape for x in self.earlybufs]), \"all earlybufs must have the same shape\"\n    assert all_same([x.shape for x in self.bufs if x not in self.earlybufs]), \"all latebufs must have the same shape\"\n    assert all_same([len(x.shape) for x in self.bufs]), \"all bufs must have the same shape size\"\n\n    # process\n    self.sts : List[ShapeTracker] = [x.st.copy() for x in self.bufs]   # create new shapetrackers inside this kernel\n    self.simplify_ones()\n    self.simplify_merge_adjacent()\n\n  def print(self):\n    buf_count = -1\n    op_count = -1\n    cache = {}\n    def print_ast(x, name=None):\n      nonlocal buf_count, op_count\n      if x not in cache:\n        if not isinstance(x, LazyOp):\n          if name is None:\n            buf_count += 1\n            name = f\"buf{buf_count}\"\n          print(f\"buf{buf_count} = {x}\")\n          cache[x] = name\n        else:\n          srcs = [print_ast(y) for y in x.src]\n          if name is None:\n            op_count += 1\n            name = f\"op{op_count}\"\n          print(f\"{name} = LazyOp({str(x.op)}, ({','.join(srcs)},), {x.arg})\")\n          cache[x] = name\n      return cache[x]\n    print_ast(self.input_ast, \"ast\")\n\n  def printbufs(self, prefix=\"\"):\n    print(f\"first_reduce: {self.first_reduce} shape_len: {self.shape_len} group_for_reduce: {self.group_for_reduce}\")\n    for i in range(len(self.sts)):\n      print(prefix, self.buftokens[i], f\"early:{'T' if i < len(self.bufs) and self.bufs[i] in self.earlybufs else 'F'}\", self.sts[i].shape, self.sts[i].views[-1].strides, len(self.sts[i].views), type(self.bufs[i]._buf))\n\n  @property\n  def shape_len(self) -> int: return len(self.sts[0].shape)\n\n  def simplify_ones(self):\n    # remove places where the shape is all ones\n    # TODO: this should be factored in to multi shape stride\n    all_ones = [all(st.shape[i]==1 for st in self.sts) for i in range(self.shape_len)]\n    # keep at least 1 one\n    if all(all_ones): all_ones[-1] = False\n    self.reshape_and_permute(lambda shape: [x for i,x in enumerate(shape) if not all_ones[i]], None)\n    # find first mismatch, don't reduce this\n    self.first_reduce = get_first_reduce([x.shape for x in self.sts])\n\n  def simplify_merge_adjacent(self):\n    shapes, strides = [x.shape for x in self.sts], [x.views[-1].strides for x in self.sts]\n\n    # merge dimensions if we can, multi get_shape_strides\n    # TODO: does this always preserve the reduce dimension, NO\n    # TODO: move this into shapetracker, with tests!\n    rets = [[(shapes[j][0], strides[j][0])] for j in range(len(shapes))]\n    for i in range(1, len(shapes[0])):\n      can_merge = []\n      for j in range(len(shapes)):\n        # TODO: added the always mergeability of 1s, is this right? if so, add to shapetracker in the 1 case\n        can_merge.append((strides[j][i] != 0 and rets[j][-1][1] == shapes[j][i]*strides[j][i]) or (strides[j][i] == 0 and rets[j][-1][1] == 0))\n      # more can merge than this\n      mergeable = all(can_merge) and i != self.first_reduce\n      for j in range(len(shapes)):\n        if mergeable:\n          rets[j][-1] = (rets[j][-1][0] * shapes[j][i], strides[j][i])\n        else:\n          rets[j].append((shapes[j][i], strides[j][i]))\n\n    for i,x in enumerate(rets): self.sts[i].reshape(tuple(y[0] for y in x))\n    self.first_reduce = get_first_reduce([x.shape for x in self.sts])\n\n  # this should be aware of the three parts to the shape\n  #  * the input/output dimensions\n  #  * the reduce dimensions\n  #  * the size outputted by each kernel\n  def reshape_and_permute(self, new_shape_fxn, axis):\n    for st in self.sts:\n      if new_shape_fxn is not None: st.reshape(tuple(new_shape_fxn(st.shape)))\n      if axis is not None: st.permute(tuple(axis))\n\n  # drops the final dimension\n  def upcast(self, allow_float4=True):\n    upcasted = [x.shape[-1] for x in self.sts if x.shape[-1] != 1]\n    assert len(upcasted) >= 1 and all_same(upcasted), f\"can't upcast mismatch {upcasted}\"\n    for i in range(len(self.bufs)):\n      st = self.sts[i]\n      if st.shape[-1] == upcasted[0]:\n        # multiview shapetrackers can slice through a float4, so don't allow them\n        can_merge = (not st.needs_valid() and len(st.views) == 1) or \"Image\" in str(type(self.bufs[i]._buf))  # TODO: terrible hack\n        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:\n          # this is an upcast to FLOAT4\n          self.buftokens[i].typ = Types.FLOAT4\n          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))\n          assert self.sts[i].offset % upcasted[0] == 0\n        else:\n          self.buftokens[i].array(upcasted[0], st.views[-1].strides[-1], len(upcasted) != len(self.sts))\n\n    # remove the last dimension\n    for st in self.sts: st.views[-1] = View(st.shape[0:-1], st.views[-1].strides[0:-1], st.views[-1].offset)\n", "path": "tinygrad/ast.py"}]}