diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py index f4fcf48212..45e5b07012 100644 --- a/test/test_uop_graph.py +++ b/test/test_uop_graph.py @@ -410,6 +410,7 @@ class TestUOpGraph(unittest.TestCase): self.assertEqual(len(uops), 4) assert_equiv_uops(uops[-1], UOp.store(glbl, idx1, val)) + @unittest.skip("this is a uop type error") def test_asserts_bad_gate(self): glbl0 = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), (), 0) idx = UOp.const(dtypes.int, 0) diff --git a/tinygrad/codegen/uopgraph.py b/tinygrad/codegen/uopgraph.py index 474f583890..2250b356d2 100644 --- a/tinygrad/codegen/uopgraph.py +++ b/tinygrad/codegen/uopgraph.py @@ -248,8 +248,8 @@ constant_folder = PatternMatcher([ (UPat(UOps.GEP, src=(UPat((UOps.ALU, UOps.CAST, UOps.BITCAST), name='alu'),), 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)), # tensor core with a 0 input is acc - *[(UPat(UOps.WMMA, src=(UPat.const(None, 0.0), UPat.var(), UPat.var("acc"))), lambda acc: acc) for i in [2, 4, 8]], - *[(UPat(UOps.WMMA, src=(UPat.var(), UPat.const(None, 0.0), UPat.var("acc"))), lambda acc: acc) for i in [2, 4, 8]], + (UPat(UOps.WMMA, src=(UPat.const(None, 0.0), UPat.var(), UPat.var("acc"))), lambda acc: acc), + (UPat(UOps.WMMA, src=(UPat.var(), UPat.const(None, 0.0), UPat.var("acc"))), lambda acc: acc), # tensor core cleanups (UPat.var("add") + UPat(UOps.WMMA, name="wmma"), lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)), @@ -328,6 +328,8 @@ constant_folder = PatternMatcher([ # ** two stage add/mul folding ** ((UPat.var("x") + UPat.cvar("c1")) + UPat.cvar("c2"), lambda x,c1,c2: x+(c1+c2)), ((UPat.var("x") * UPat.cvar("c1")) * UPat.cvar("c2"), lambda x,c1,c2: x*(c1*c2)), + ((UPat.var("x") & UPat.cvar("c1")) & UPat.cvar("c2"), lambda x,c1,c2: x&(c1&c2)), + ((UPat.var("x") | UPat.cvar("c1")) | UPat.cvar("c2"), lambda x,c1,c2: x|(c1|c2)), # *** rules from symbolic *** # ** lt ** # c0*x