diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index 9053a85e6a..5dda898f0c 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -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]]: diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 3a8d7a0583..ddc3884932 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -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):