github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

simple block barrier (#11341)

Browse Source 
* simple block barrier

* simple
This commit is contained in:
George Hotz
2025-07-23 10:14:11 -07:00
committed by GitHub
parent 86e7504111
commit ae07a93814
1 changed files with 3 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tinygrad/codegen/linearize.py
View File

@@ -86,21 +86,13 @@ class BlockContext:
ctx.child_count[s] += 1
this_block_ctx += ctx.last_ctx(s)
# save the block ctx
ctx.block_ctxs[u] = _sort_ctx(this_block_ctx)
# save the block ctx. SINK never has anything
ctx.block_ctxs[u] = _sort_ctx(this_block_ctx) if u.op is not Ops.SINK else ()
# RANGE/IF add to the next ctx
# STORE/ASSIGN subtract from the next ctx
if u.op in {Ops.RANGE, Ops.IF}: ctx.child_ctxs[u] = _sort_ctx(ctx.block_ctxs[u] + (u,))
elif u.op is Ops.STORE:
if len(definereg:=[x for x in u.src[0].toposort() if x.op is Ops.DEFINE_REG]):
# old assign logic
ctx.child_ctxs[u] = tuple([y for y in ctx.last_ctx(u.src[1]) if y not in definereg[0].src[1:]])
elif any(x.op is Ops.DEFINE_LOCAL for x in u.src[0].toposort()):
# deal with non-reduce locals. probably wrong
idx_context, store_context = ctx.last_ctx(u.src[0]), ctx.last_ctx(u.src[1])
ctx.child_ctxs[u] = tuple([y for y in store_context if y not in idx_context and y.op is Ops.RANGE])
else: ctx.child_ctxs[u] = ()
elif u.op is Ops.STORE: ctx.child_ctxs[u] = tuple([y for y in ctx.block_ctxs[u] if y not in u.src])
return ctx
# ***** make blocks *****