weak dtypes in dtype_from_uop [PR] (#17032)

* weak dtypes in dtype_from_uop [PR]

* no weak in spec_program

* weak const fold tests

test/null/test_const_folding.py

@@ -1,6 +1,6 @@
 import unittest, itertools, math
 from tinygrad import Tensor, dtypes, Context
-from tinygrad.dtype import DType, ConstType
+from tinygrad.dtype import DType, ConstType, Invalid
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
@@ -34,6 +34,24 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
     x = x.clip(0, 1).realize()
     _check_ast_count(1, x.neg())
 
+class TestWeakConstFolding(unittest.TestCase):
+  def test_weakint_math(self):
+    out = (UOp.const(dtypes.weakint, 2**40) + UOp.const(dtypes.weakint, 2**40)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41))
+
+  def test_float_unaries(self):
+    for dtype in (dtypes.weakint, dtypes.weakfloat):
+      for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
+        out = UOp.const(dtype, 4).alu(op).simplify()
+        self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))
+
+  def test_weakfloat_math(self):
+    out = (UOp.const(dtypes.weakfloat, 1.25) + UOp.const(dtypes.weakfloat, 2.5)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakfloat, 3.75))
+
+  def test_invalid_poison(self):
+    self.assertIs(UOp.const(dtypes.weakint, Invalid).alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid)
+
 class TestBinaryOpsConstFolding(unittest.TestCase):
   def test_add_literal_zero(self):
     _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)

test/null/test_uops.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
 from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
-from tinygrad.uop.spec import spec_shared
+from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym
 from test.helpers import eval_uop, to_uops_list
 
@@ -26,6 +26,27 @@ class TestDTypeFromUOp(unittest.TestCase):
     idx = UOp.range(4, 0)
     self.assertEqual(idx.valid(idx < 4).dtype, dtypes.index)
 
+  def test_const_dtype_from_value(self):
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), True), dtypes.bool)
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), 3), dtypes.weakint)
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), ConstFloat(3.0)), dtypes.weakfloat)
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), Invalid), dtypes.bool)
+    self.assertRaises(TypeError, dtype_from_uop, Ops.CONST, (), (1, 2))
+
+  @Context(SPEC=2)
+  def test_const_default_dtype_is_derived(self):
+    self.assertEqual(UOp(Ops.CONST, arg=3).dtype, dtypes.weakint)
+    self.assertEqual(UOp(Ops.CONST, arg=ConstFloat(3.0)).dtype, dtypes.weakfloat)
+    self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool)
+    self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool)
+    # an explicit (strong) const dtype is legal until the field is removed
+    self.assertEqual(UOp.const(dtypes.int32, 3).dtype, dtypes.int32)
+
+  def test_weak_dtype_rejected_by_program_spec(self):
+    for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
+      with self.assertRaises(RuntimeError): type_verify(UOp.const(weak, value).sink(), spec_program)
+      type_verify(UOp.const(concrete, value).sink(), spec_program)
+
 class TestSafeCast(unittest.TestCase):
   def test_cast_folds(self):
     a = UOp.variable("a", 1, 10, dtype=dtypes.int32)

tinygrad/uop/ops.py

@@ -156,8 +156,12 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
       assert isinstance(arg, DType), f"CAST/BITCAST arg must be DType, got {arg}"
       return arg
     case Ops.CONST:
-      # TODO: need const refactor to bool/weakint/weakfloat
-      return None
+      # derived from the value. order matters: bool is an int subclass, ConstFloat is a float subclass
+      if isinstance(arg, InvalidType): return dtypes.bool  # Invalid is the lattice bottom, typed by its consumer
+      if isinstance(arg, bool): return dtypes.bool
+      if isinstance(arg, int): return dtypes.weakint
+      if isinstance(arg, float): return dtypes.weakfloat
+      raise TypeError(f"no dtype for CONST with arg {arg}")
   if op in GroupOp.Unary: return src[0].dtype
   # NOTE: CMPLT, CMPNE, CMPEQ, WHERE, SHL, SHR are handled above
   if op in GroupOp.Broadcastable:
@@ -171,7 +175,8 @@ class UOpMetaClass(type):
   def __call__(cls, op:Ops, dtype:DType|None=None, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
                metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None):
     if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void
-    if SPEC == 2 and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype:
+    # CONST derives its dtype by value only when the constructor omits one; an explicit (strong) const dtype is legal until the field is removed
+    if SPEC == 2 and op is not Ops.CONST and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype:
       raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}")
     if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret
     UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key))

tinygrad/uop/spec.py

@@ -195,8 +195,8 @@ spec_tensor = PatternMatcher([
 
 # these ops can exist in programs but not the tensor spec. example: LOAD
 spec_program = PatternMatcher([
-  # index is not allowed in programs
-  (UPat(GroupOp.All, dtypes.index), lambda: False),
+  # index and weak dtypes are not allowed in programs
+  (UPat(GroupOp.All, (dtypes.index, dtypes.weakint, dtypes.weakfloat)), lambda: False),
 
   # allow special SHRINK
   (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST))), lambda: True),