diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 072c8fe11b..a3997e4f68 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -1,7 +1,7 @@ name: Unit Tests env: # increment this when downloads substantially change to avoid the internet - DOWNLOAD_CACHE_VERSION: '10' + DOWNLOAD_CACHE_VERSION: '11' PYTHON_CACHE_VERSION: '2' CAPTURE_PROCESS_REPLAY: 1 GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 6918062d9a..48bc9461f0 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -242,18 +242,18 @@ def gep_through_wmma(gep:UOp, wmma:UOp): gep_pushing = PatternMatcher([ # GEP/VECTORIZE, GEP/GEP, GEP/CONST, GEP/VCONST - (UPat(Ops.GEP, src=(UPat(Ops.GEP, name='g2'),), name='g1'), + (UPat(Ops.GEP, name='g2').f(Ops.GEP, name='g1'), lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(len(g1.arg))))), - (UPat(Ops.GEP, src=(UPat(Ops.VECTORIZE, name="vec"),), name="gep"), + (UPat(Ops.VECTORIZE, name='vec').f(Ops.GEP, name='gep'), lambda gep, vec: UOp(Ops.VECTORIZE, gep.dtype, tuple(vec.src[i] for i in gep.arg)) if len(gep.arg) > 1 else vec.src[gep.arg[0]]), - (UPat(Ops.GEP, src=(UPat.cvar("c", vec=False),), name="gep"), lambda gep, c: gep.const_like(c.arg)), - (UPat(Ops.GEP, src=(UPat(Ops.VCONST, name="c"),), name="gep"), lambda gep, c: gep.const_like(tuple(c.arg[x] for x in gep.arg))), + (UPat.cvar("c", vec=False).f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(c.arg)), + (UPat(Ops.VCONST, name="c").f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(tuple(c.arg[x] for x in gep.arg))), # GEP on void is skipped (UPat(Ops.GEP, src=(UPat(dtype=dtypes.void, name="x"),)), lambda x: x), # GEP in order is removed (UPat(Ops.GEP, name="g"), lambda g: g.src[0] if not isinstance(g.dtype, PtrDType) and g.arg == tuple(range(g.src[0].dtype.count)) else None), # push all GEPs through ALUs (fix arange stuff) - (UPat(Ops.GEP, src=(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu'),), name='gep'), + (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, name='gep'), lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg) \ if not isinstance(gep.dtype, PtrDType) else None), # CAT can't be rendered. it's a VECTORIZE on vectors, we expand to a single VECTORIZEs with GEPs (TODO: move this later) @@ -262,7 +262,7 @@ gep_pushing = PatternMatcher([ # VECTORIZE on same GEP (UPat(Ops.VECTORIZE, name="v", src=UPat(Ops.GEP, src=(UPat.var("x"),))), lambda v,x: x.gep(tuple(get_single_element(i.arg) for i in v.src))), # push some GEPs through WMMAs - (UPat(Ops.GEP, src=(UPat(Ops.WMMA, name="wmma"),), name="gep"), gep_through_wmma), + (UPat(Ops.WMMA, name="wmma").f(Ops.GEP, name="gep"), gep_through_wmma), ]) commutative = PatternMatcher([