diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index c3c7657805..b81204342b 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -82,13 +82,13 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC supported_ops = tuple(opts.code_for_op.keys()) extra_matcher = opts.extra_matcher if opts.extra_matcher is not None else PatternMatcher([]) + # optional pre matcher + if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher")) + # decompositions pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, _TRANSCENDENTAL>=2) ret.append(RewriteStep(pm_decomp, name="decompositions")) - # optional pre matcher - if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher")) - # final rules for the renderer (without sym) pm_final_rewrite = symbolic_simple+pm_render+extra_matcher ret.append(RewriteStep(pm_final_rewrite, lambda _: opts.device, name="final rewrite")) diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index 4063a23c76..513d03018a 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -28,14 +28,17 @@ def is_packed(dt:DType, odt:DType|None = None) -> bool: if odt is None: odt = dt return dt.itemsize < 4 and dt.base != dtypes.half and (not isinstance(odt, PtrDType) or odt.addrspace != AddrSpace.REG) -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)), +wgsl_pack_load_store = PatternMatcher([ (UPat.load(UPat.var("b"), UPat.cvar("c"), name="l"), lambda l,b,c: packed_load(l,b,l.dtype,c.cast(dtypes.uint32)) if is_packed(l.dtype, b.dtype) else None), (UPat.load(UPat.var("b"), name='l', allow_any_len=True), lambda l,b: packed_load(l, b, l.dtype) if is_packed(l.dtype, b.dtype) else None), (UPat.store(UPat.var("bidx"), UPat.var("var"), allow_any_len=True), lambda bidx,var: packed_store(bidx,var) if is_packed(var.dtype, bidx.dtype) else None), +]) + +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)), (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None), ]) + extra_pm @@ -45,6 +48,7 @@ class WGSLRenderer(CStyleLanguage): global_max = (65535, 65535, 65535) local_max = (256, 256, 64) code_for_workitem = {"g": lambda x: f"i32(gindex.{'xyz'[int(x)]})", "l": lambda x: f"i32(lindex.{'xyz'[int(x)]})"} + pre_matcher = wgsl_pack_load_store extra_matcher = wgsl_matcher supports_float4 = False barrier = "workgroupBarrier();"