github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-26 15:28:10 -05:00
Code Issues Packages Projects Releases Wiki Activity

minor Program post_init size cleanup [run_process_replay] (#6415)

Browse Source
This commit is contained in:
chenyu
2024-09-08 23:41:27 -04:00
committed by GitHub
parent 442150a8df
commit 047ab7d256
1 changed files with 3 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tinygrad/renderer/__init__.py
View File

@@ -40,12 +40,9 @@ class Program:
if u.op is UOps.SPECIAL:
# NOTE: you have to set local_size and global_size to the base [1,1,1] outside this
if u.arg[0][0] == 'i': self.local_size = None
if u.arg[0][0] == 'l':
assert self.local_size is not None
self.local_size[int(u.arg[0][-1])] = u.arg[1]
else:
assert self.global_size is not None
self.global_size[int(u.arg[0][-1])] = u.arg[1]
special_size = self.local_size if u.arg[0][0] == 'l' else self.global_size
assert special_size is not None
special_size[int(u.arg[0][-1])] = u.arg[1]
self.vars = sorted(self.vars, key=lambda v: v.expr)
self.outs = sorted(dedup(self.outs))
self._ran_post_init = True