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

fix add_raw_barrier [pr] (#17278)

Browse Source 
* fix add_raw_barrier [pr]

a kernel might need multiple barriers

* smaller?
This commit is contained in:
chenyu
2026-07-29 09:30:00 -07:00
committed by GitHub
parent 6ea7d366fa
commit 3803f1583b
2 changed files with 13 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_linearizer.py
+11
View File
@@ -13,6 +13,7 @@ from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.cstyle import CUDARenderer
from tinygrad.renderer.isa import ISARenderer
from test.helpers import replace_opts
from test.backend.test_softmax_fusion import single_kernel_softmax
MOCKGPU = DEV.interface.startswith("MOCK")
from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
@@ -392,6 +393,16 @@ class TestLinearizer(unittest.TestCase):
# the global store doesn't change
assert stores[1].src[1].dtype == dtypes.float
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
def test_two_grouped_stores_local(self):
# GROUP on both reduces puts two LOCAL buffers in one kernel, and the store to each needs its own barrier
a = Tensor.rand(32, 32).realize()
opts = [Opt(OptOps.GROUP, 1, 4), Opt(OptOps.GROUP, 2, 4)]
ast = helper_linearizer_opt(single_kernel_softmax(a), [opts])
uops = to_program(replace_opts(ast, opts), renderer=Device[Device.DEFAULT].renderer).src[1].src
self.assertEqual(len([u for u in uops if u.op is Ops.BARRIER]), 2)
# *** helpers ***
def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
tinygrad/codegen/__init__.py
+2 -2
View File
@@ -261,8 +261,8 @@ def add_raw_barrier(after:UOp):
# loads from a LOCAL buffer that depend (via AFTER) on stores to LOCAL memory need a workgroup barrier
if after.addrspace is not AddrSpace.LOCAL: return None
# one toposort over all the deps
deps = UOp.sink(*after.src[1:]).backward_slice
if not any(_is_local_store(x) for x in deps) or any(x.op is Ops.BARRIER for x in deps): return None
deps = UOp.sink(*after.src[1:]).toposort(gate=lambda x: x.op is not Ops.BARRIER)
if not any(_is_local_store(x) for x in deps): return None
return after.src[0].after(UOp(Ops.BARRIER, src=after.src[1:]))
def add_war_barrier(end:UOp):