geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

reorder UOps.DEFINE_VAR in runtime [run_process_replay] (#5659)

Browse Source 
prep rewrite SPECIAL using DEFINE_VAR
This commit is contained in:
chenyu
2024-07-23 14:32:10 -04:00
committed by GitHub
parent 199b3bf02b
commit fdc72ba102
3 changed files with 11 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/runtime/ops_python.py
+2 -4
View File
@@ -84,10 +84,8 @@ class PythonProgram:
elif uop is UOps.DEFINE_VAR:
ul[i] = [pvals.pop(0)] * warp_size
elif uop is UOps.SPECIAL:
if arg[1][0] == 'g':
ul[i] = [idxs[2-arg[0]]] * warp_size
elif arg[1][0] == 'l':
ul[i] = [x[2-arg[0]] for x in warp]
if arg[1][0] == 'g': ul[i] = [idxs[2-arg[0]]] * warp_size
elif arg[1][0] == 'l': ul[i] = [x[2-arg[0]] for x in warp]
elif uop is UOps.CONST:
ul[i] = [[arg] * warp_size for _ in range(dtype.count)] if dtype.count > 1 else [arg] * warp_size
elif uop is UOps.DEFINE_ACC: