diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index e9d26d0113..a10d13ea8b 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -307,8 +307,8 @@ class TestRecurse(unittest.TestCase): def test_inf_loop(self): a = UOp.const(3) pm = PatternMatcher([ - (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), - (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), + (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), + (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) with self.assertRaises(RuntimeError): graph_rewrite(a, pm) @@ -316,8 +316,8 @@ class TestRecurse(unittest.TestCase): def test_inf_loop_bottom_up(self): a = UOp.const(3) pm = PatternMatcher([ - (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), - (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), + (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), + (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) with self.assertRaises(RuntimeError): graph_rewrite(a, pm, bottom_up=True) @@ -378,8 +378,8 @@ class TestWalkRewrite(unittest.TestCase): """A bouncing pattern applies once and stops instead of looping.""" a = UOp.const(3) pm = PatternMatcher([ - (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), - (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), + (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), + (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) with self.assertRaises(RuntimeError): graph_rewrite(a, pm, bottom_up=True) @@ -456,8 +456,8 @@ class TestWalkRewrite(unittest.TestCase): """Bottom-up walk also applies once per node, no fixed-point iteration.""" a = UOp.const(3) pm = PatternMatcher([ - (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), - (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), + (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), + (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) ret = graph_rewrite(a, pm, bottom_up=True, walk=True) self.assertIs(ret, UOp.const(4)) @@ -511,7 +511,7 @@ class TestWalkRewrite(unittest.TestCase): def bpm_match(ctx, x): ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm")) # rewrite const(1) -> const(10), short-circuiting its subtree - if x.op is Ops.CONST and x.val == 1: return x.replace(arg=10) + if x.op is Ops.CONST and x.val == 1: return UOp.const(10, x.dtype) return None def pm_match(ctx, x): ctx.append((x.val if x.op is Ops.CONST else x.op, "pm")) diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index e4ce98c0db..e889df8611 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -593,7 +593,7 @@ class TestUOpTags(unittest.TestCase): def test_inc_by_one(self): g = UOp.const(1) + UOp.const(1) assert g.ssimplify() == 2 - pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.val+1, tag=1) if x.tag is None else None)]) + pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: UOp.const(x.val+1, x.dtype).rtag(1) if x.tag is None else None)]) pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)]) g = graph_rewrite(g, pm_plus_1) assert g.ssimplify() == 4 diff --git a/test/null/test_uops.py b/test/null/test_uops.py index 09f6e9a55a..247c611250 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -126,6 +126,12 @@ class TestConstFloatEq(unittest.TestCase): self.assertFalse(nan == Invalid) self.assertTrue(nan != Invalid) # __ne__ must defer to the reflected eq, not swallow NotImplemented + def test_invalid_eq_defers_to_reflected(self): + class HoldsInvalid: # a carrier that knows it holds Invalid. returning False for foreign types would silence its eq + def __eq__(self, other): return other is Invalid + self.assertTrue(Invalid == HoldsInvalid()) + self.assertFalse(Invalid != HoldsInvalid()) + def test_matchers_agree_on_nan(self): n = UOp.const(math.nan, dtypes.float32) for compiled in (False, True): @@ -447,7 +453,7 @@ class TestUPatHelpers(unittest.TestCase): class TestUopsObject(unittest.TestCase): def test_timing(self): - with Timing("create 10k uops:"): ret = [UOp(Ops.CONST, dtypes.int, arg=10000000+i) for i in range(10000)] + with Timing("create 10k uops:"): ret = [UOp.const(10000000+i, dtypes.int) for i in range(10000)] assert len(ret) == 10000 def test_nested(self): diff --git a/test/null/test_uops_stats.py b/test/null/test_uops_stats.py index 7bef6acb1b..edab948072 100644 --- a/test/null/test_uops_stats.py +++ b/test/null/test_uops_stats.py @@ -147,21 +147,21 @@ class TestUOpsStats(unittest.TestCase): #MULACC should have the same stats as MUL + ADD def test_mulacc(self): globl = UOp.param(0, dtypes.int, (3,)) - o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1) - o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2) + o1 = UOp.const(1, dtypes.int) + o2 = UOp.const(2, dtypes.int) u1 = globl.index(o1) u2 = globl.index(o2) - u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3) + u3 = UOp.const(3, dtypes.int) u4 = UOp(Ops.MUL, src=(u1,u2)) u5 = UOp(Ops.ADD, src=(u4,u3)) uops = tuple(u5.toposort()) globl = UOp.param(0, dtypes.int, (3,)) - o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1) - o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2) + o1 = UOp.const(1, dtypes.int) + o2 = UOp.const(2, dtypes.int) u1 = globl.index(o1) u2 = globl.index(o2) - u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3) + u3 = UOp.const(3, dtypes.int) u4 = UOp(Ops.MULACC, src=(u1,u2,u3)) uops_fma = tuple(u4.toposort()) diff --git a/test/null/test_viz.py b/test/null/test_viz.py index fad98fc4a4..954a4f2252 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -97,7 +97,7 @@ class TestViz(unittest.TestCase): # VIZ tracks rewrites up to and including the error def count_3(x:UOp): assert x.val <= 3 - return x.replace(arg=x.val+1) + return UOp.const(x.val+1, x.dtype) err_pm = PatternMatcher([(UPat.cvar("x"), count_3),]) a = UOp.const(1) with save_viz() as viz: @@ -202,8 +202,8 @@ class TestViz(unittest.TestCase): a = UOp.const(3) b = UOp.const(4) pm = PatternMatcher([ - (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), - (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), + (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), + (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) with save_viz() as viz: # use smaller stack limit for faster test (default is 250000) @@ -224,7 +224,7 @@ class TestViz(unittest.TestCase): list(viz.get_details(0, 0)) def test_enter_calls_rewrite(self): - pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4))]) + pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype))]) with save_viz() as viz: inner = UOp.const(3) call = UOp(Ops.CALL, src=(UOp(Ops.SINK, src=(inner,)),)) diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index d682aad8f5..cbb324ace0 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -178,7 +178,9 @@ pm_float_decomp = PatternMatcher([ f2f(x.bitcast(f2f_dt[ctx[0]]), ctx[0], ctx[1]) if bc.dtype == ctx[0] else None), (UPat(Ops.CAST, dtypes.floats, src=(UPat.var("val"),), name="x"), lambda ctx,x,val: f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype == ctx[0] else None), - (UPat(GroupOp.All-{Ops.BITCAST}, dtypes.floats, name="x"), lambda ctx,x: + # a CONST has no srcs to cast, it restates its value at the emulating dtype + (UPat(Ops.CONST, dtypes.floats, name="x"), lambda ctx,x: UOp.const(x.val, ctx[1]) if x.dtype == ctx[0] else None), + (UPat(GroupOp.All-GroupOp.Defines-{Ops.CAST, Ops.BITCAST, Ops.CONST}, dtypes.floats, name="x"), lambda ctx,x: x.replace(dtype=ctx[1], src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src)) if x.dtype == ctx[0] else None), (UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), name='st'), lambda ctx,st,idx,val: diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 9272b7e617..4558aeaa46 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -27,7 +27,7 @@ class InvalidType: def __new__(cls): if cls._instance is None: cls._instance = object.__new__(cls) return cls._instance - def __eq__(self, other): return self is other + def __eq__(self, other): return self is other if isinstance(other, InvalidType) else NotImplemented # foreign types get the reflected eq def __hash__(self): return id(self) def __repr__(self): return "Invalid" def __reduce__(self): return (InvalidType, ()) # unpickle returns the singleton diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 0a35d6759d..1fcd8a4dba 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -54,7 +54,7 @@ def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str| return freqs.cos().cat(freqs.sin(), dim=-1).clone(device) class ExpertWeights: - """Like nn.Linear but with num_experts dimension. Weight shape: (num_experts, out_features, in_features).""" + """Like Linear but with num_experts dimension. Weight shape: (num_experts, out_features, in_features).""" def __init__(self, num_experts:int, in_features:int, out_features:int): self.weight = Tensor.zeros(num_experts, out_features, in_features) def __call__(self, sel:Tensor, x:Tensor) -> Tensor: @@ -123,15 +123,15 @@ class FFNBlock: self.pending_state:tuple[Tensor, Tensor]|None = None self.attn_norm, self.ffn_norm = nn.RMSNorm(config.dim, config.norm_eps), nn.RMSNorm(config.dim, config.norm_eps) if config.num_experts > 0: - self.ffn_gate_inp = nn.Linear(config.dim, config.num_experts, bias=False) + self.ffn_gate_inp = Linear(config.dim, config.num_experts, bias=False) # router if config.expert_bias: self.exp_probs_b = {"bias": Tensor.zeros(config.num_experts)} self.ffn_gate_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim) self.ffn_up_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim) self.ffn_down_exps = ExpertWeights(config.num_experts, config.hidden_dim, config.dim) if config.shared_expert_dim > 0: - self.ffn_gate_shexp = nn.Linear(config.dim, config.shared_expert_dim, bias=False) - self.ffn_up_shexp = nn.Linear(config.dim, config.shared_expert_dim, bias=False) - self.ffn_down_shexp = nn.Linear(config.shared_expert_dim, config.dim, bias=False) + self.ffn_gate_shexp = Linear(config.dim, config.shared_expert_dim, bias=False) + self.ffn_up_shexp = Linear(config.dim, config.shared_expert_dim, bias=False) + self.ffn_down_shexp = Linear(config.shared_expert_dim, config.dim, bias=False) if config.shared_expert_gate: self.ffn_gate_inp_shexp = {"weight": Tensor.zeros(config.dim)} else: self.ffn_gate, self.ffn_up = Linear(config.dim, config.hidden_dim, bias=False), Linear(config.dim, config.hidden_dim, bias=False) @@ -283,16 +283,16 @@ class MLATransformerBlock(FFNBlock): super().__init__(config) qk_nope_head_dim = config.head_dim - config.rope_dim if config.q_lora_rank > 0: - self.attn_q_a = nn.Linear(config.dim, config.q_lora_rank, bias=False) + self.attn_q_a = Linear(config.dim, config.q_lora_rank, bias=False) self.attn_q_a_norm = nn.RMSNorm(config.q_lora_rank, config.norm_eps) - self.attn_q_b = nn.Linear(config.q_lora_rank, config.n_heads * config.head_dim, bias=False) + self.attn_q_b = Linear(config.q_lora_rank, config.n_heads * config.head_dim, bias=False) else: - self.attn_q = nn.Linear(config.dim, config.n_heads * config.head_dim, bias=False) - self.attn_kv_a_mqa = nn.Linear(config.dim, config.kv_lora_rank + config.rope_dim, bias=False) + self.attn_q = Linear(config.dim, config.n_heads * config.head_dim, bias=False) + self.attn_kv_a_mqa = Linear(config.dim, config.kv_lora_rank + config.rope_dim, bias=False) self.attn_kv_a_norm = nn.RMSNorm(config.kv_lora_rank, config.norm_eps) self.attn_k_b = {"weight": Tensor.zeros(config.n_heads, config.kv_lora_rank, qk_nope_head_dim)} self.attn_v_b = {"weight": Tensor.zeros(config.n_heads, config.v_head_dim, config.kv_lora_rank)} - self.attn_output = nn.Linear(config.n_heads * config.v_head_dim, config.dim, bias=False) + self.attn_output = Linear(config.n_heads * config.v_head_dim, config.dim, bias=False) def _attention(self, x:Tensor, start_pos:int|UOp, use_flash:bool=False, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None) -> Tensor: diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 67129a675d..0b20d84b24 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1764,7 +1764,7 @@ def lower_weak_node(u:UOp) -> UOp|None: else unwrap(dtype_from_uop(u.op, src, u.arg))) return u.replace(dtype=None, src=src[:start]+tuple(s.cast(dt) for s in src[start:])).cast(u.dtype) pm_lower_weak = PatternMatcher([ - (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype)), + (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)), # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default. # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs) (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"),