Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Multireduce Kernels: Allow IF blocks to terminate early #4457

Closed
wants to merge 5 commits into from
Closed

Multireduce Kernels: Allow IF blocks to terminate early #4457

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
962b44c
allow endifs to be inserted before the end of the graph
0xtimmy May 6, 2024
bbce926
Merge remote-tracking branch 'upstream/master' into early-terminate-if
0xtimmy May 6, 2024
18641c9
find add ENDIF before next BARRIER
0xtimmy May 8, 2024
f1bbd71
removing tests with manual ENDIF + linters
0xtimmy May 8, 2024
b288a5c
specifically the next barrier aftr the store of the local result
0xtimmy May 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
8 changes: 8 additions & 0 deletions test/test_uop_graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,5 +46,13 @@ def test_const_cast(self):
self.assertEqual(out.uop, UOps.CONST)
self.assertEqual(out.arg, 0)

def test_early_endif(self):
g = UOpGraph()
g.add(UOps.IF, vin=(g.add(UOps.CONST, dtypes.bool, arg=True),), cachable=False)
g.add(UOps.CONST, dtypes.int, arg=0)
g.add_ends()
self.assertEqual(len([x for x in g.uops if x.uop is UOps.ENDIF]), 1, "UOpGraph.add_ends() should not add any extra ENDIFs")
self.assertEqual(g.uops[-1].uop, UOps.ENDIF, "UOpGraph.add_ends() should add ENDIF to the end of the graph")

if __name__ == '__main__':
unittest.main(verbosity=2)
9 changes: 5 additions & 4 deletions tinygrad/codegen/uops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,13 +193,13 @@ def type_verify(self):
assert vin[0].dtype == dtypes.bool, f"{arg} selector dtype mismatch {vin[0].dtype=} != {dtypes.bool}"
assert dtype == vin[1].dtype == vin[2].dtype, f"{arg} choice dtype mismatch {dtype=} != {vin[1].dtype=} != {vin[2].dtype=}"

def get_recursive_children(self, x:UOp) -> Set[UOp]:
def get_recursive_children(self, x:UOp, with_phi=False) -> Set[UOp]:
deps = set([x])
ssize = 0
while ssize != len(deps):
ssize = len(deps)
for u in self.uops:
if len(deps.intersection([x for x in u.vin if x.uop is not UOps.PHI])):
if len(deps.intersection([x for x in u.vin if with_phi or x.uop is not UOps.PHI])):
deps.add(u)
return deps

Expand All @@ -210,8 +210,9 @@ def add_ends(self):
insert_before = self.uops.index(sorted(list(self.get_recursive_children(u)), key=self.uops.index)[-1])+1
self.add(UOps.ENDLOOP, None, (u,), cachable=False, insert_before=insert_before)
elif u.uop is UOps.IF:
# END any if statements at the end of the uops
self.add(UOps.ENDIF, None, (u,), cachable=False)
# add END of if after the barrier of the store of the result
barriers = (sorted([self.uops.index(x) for x in list(self.get_recursive_children(u, with_phi=True)) if x.uop is UOps.BARRIER]) + [None])[0]
self.add(UOps.ENDIF, None, (u,), cachable=False, insert_before=barriers)

def fix_loop_scope(self, get_recursive_parents:Callable[..., Set[UOp]]):
loop_stack: List[List[UOp]] = [[]]
Expand Down