diff --git a/test/backend/test_renderer_failures.py b/test/backend/test_renderer_failures.py index d1f7a145fa..403a938cdc 100644 --- a/test/backend/test_renderer_failures.py +++ b/test/backend/test_renderer_failures.py @@ -77,6 +77,14 @@ class TestCStyleFailures(unittest.TestCase): @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer") class TestWGSLFailures(unittest.TestCase): + def test_folded_packed_store(self): + b = UOp.param(0, dtypes.char, (4,)) + idx = b.index(UOp.const(0).cast(dtypes.int)) + store = UOp.store(idx, UOp.load(idx, dtype=dtypes.uint32) & UOp.const(0xffffff00).cast(dtypes.uint32)) + src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort()) + self.assertIn("atomicAnd(&data0_4[0],4294967040u);", src) + self.assertNotIn("atomicAdd", src) + def test_multiply_infinity(self): # multiplying a positive constant by infinity should return infinity # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index db3bd00ed8..0f27eb0e16 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -149,6 +149,13 @@ class TestSymbolic(unittest.TestCase): def test_xor_0(self): self.helper_test_variable(Variable("a", 0, 8, dtypes.int) ^ 0, 0, 8, "a", test_z3=False) + def test_or_0(self): + self.helper_test_variable(Variable("a", 0, 8, dtypes.int) | 0, 0, 8, "a", test_z3=False) + + def test_shift_0(self): + self.helper_test_variable(Variable("a", 0, 8, dtypes.int) << 0, 0, 8, "a") + self.helper_test_variable(Variable("a", 0, 8, dtypes.int) >> 0, 0, 8, "a") + def test_xor_self_inverse(self): self.helper_test_variable((Variable("a", 0, 8, dtypes.int) ^ 5) ^ 5, 0, 8, "a", test_z3=False) diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index f032cf435a..6d34cc70c9 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -38,6 +38,9 @@ def is_nan(a): bs, (exp, mant) = a.dtype.bitsize, dtypes.finfo(a.dtype) return (a.bitcast(getattr(dtypes, f"uint{bs}")) & ((1 << (bs - 1)) - 1)) > (((1 << exp) - 1) << mant) +# the read-modify-write packed_store emits: a load of the very index being stored to, masked (a gated store loads with 3 srcs) +packed_rmw = UPat(Ops.LOAD, src=(UPat.var("b"),), allow_any_len=True) & UPat.var("wmask") + wgsl_matcher = PatternMatcher([ (UPat((Ops.CMPLT, Ops.XOR), src=(UPat(name="a", dtype=dtypes.bool), UPat.var("b")), name="c"), lambda a,b,c: a.cast(dtypes.int).alu(c.op, b.cast(dtypes.int)).cast(dtypes.bool)), @@ -82,10 +85,10 @@ class WGSLRenderer(CStyleLanguage): (UPat.load(UPat.var("b"), UPat.var("v"), UPat.var("gate")), lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b], b.src[0])}, {ctx[gate]})"), (UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b)), - (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v:\ - # (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1] - f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b) \ - else f"{ctx[b]} = {ctx[v]};"), + # packed_store writes (load & wmask) | new_v: atomicAnd clears the field, atomicAdd sets it. new_v is gone when it is 0 + (UPat.store(UPat.var("b"), UPat.any(packed_rmw, packed_rmw | UPat.var("nv"))), lambda ctx,b,wmask,nv=None: + f"atomicAnd(&{ctx[b]},{ctx[wmask]});"+(f"\n atomicAdd(&{ctx[b]},{ctx[nv]});" if nv is not None else "") if is_packed(b) else None), + (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v: f"{ctx[b]} = {ctx[v]};"), (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))), lambda ctx,b,idx: f"{ctx[b]}[{strip_parens(ctx[idx]) if idx.arg is Ops.ADD else ctx[idx]}]"), ]) + base_rewrite diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 3e4c1fd0ce..7f48a18ce2 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -108,9 +108,9 @@ def fold_const_where(gate:UOp, c0:UOp, c1:UOp, w:UOp) -> UOp: symbolic_simple = pm_data_invalid + PatternMatcher([ # ** self folding ** - (UPat.var("x") + 0, lambda x: x), # x+0 -> x + (UPat({Ops.ADD, Ops.XOR, Ops.OR}, src=[UPat.var("x"), UPat.const(0)]), lambda x: x), # x+0 / x^0 / x|0 -> x + (UPat({Ops.SHL, Ops.SHR}, src=(UPat.var("x"), UPat.const(0))), lambda x: x), # x<<0 / x>>0 -> x (UPat.var("x") * 1, lambda x: x), # x*1 -> x - (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) ^ 0, lambda x: x), # x^0 -> x (UPat.var("x") // UPat.var("x"), lambda x: x.const_like(1)), # x//x -> 1 (UPat.var("x") // 1, lambda x: x), # x//1 -> x (UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x