Revert "don't cast weak in _broadcasted [pr] (#17408)"

This reverts commit b45058b5ec.

test/unit/test_dtype_weak.py

@@ -62,15 +62,6 @@ class TestWeakPromotion(unittest.TestCase):
     self.assertEqual((x._uop.base.op, x._uop.base.val, x.dtype, x.shape, y.dtype),
                      (Ops.CONST, 1, dtypes.weakfloat, (1,), dtypes.float32))
 
-  def test_broadcasted_keeps_expression_weak(self):
-    # a weak EXPRESSION is not a const, so it survives as itself: the lub lives on the parent, not on a cast of the operand
-    x, y = Tensor([1], dtype=dtypes.int8)._broadcasted(Tensor(2) * 3)
-    self.assertEqual((y._uop.base.op, y.dtype, x.dtype), (Ops.MUL, dtypes.weakint, dtypes.int8))
-    self.assertIs((x + y).dtype, dtypes.int8)
-    # only the KIND is lifted, and a cast is the only way an expression can state it
-    x, y = Tensor([1.0], dtype=dtypes.float32)._broadcasted(Tensor(2) * 3)
-    self.assertEqual((y._uop.base.op, y.dtype), (Ops.CAST, dtypes.weakfloat))
-
   def test_uop_scalar_const_lifts_kind(self):
     for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
                                                  (dtypes.int32, 1, dtypes.int32, dtypes.weakint),

tinygrad/codegen/simplify.py

@@ -142,7 +142,7 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp
     sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
     if not no_range(sink): return None
     u = sink.substitute({v:k for k,v in replaces.items()})
-  return u.cast(red.dtype)
+  return u
 
 def reduce_load_collapse(red:UOp, u:UOp) -> UOp|None: return reduce_collapse(red, u, pm=pm_reduce_load_collapse)
 

tinygrad/mixin/elementwise.py

@@ -22,12 +22,11 @@ class ElementwiseMixin(CreationMixin):
     y = self.ufix(y)
     x, y = (self, y) if not reverse else (y, self)
     out_dtype = least_upper_dtype(x.dtype, y.dtype)
-    # a weak stays weak, might lift weakint -> weakfloat
+    # keep weak CONST weak, might lift weakint -> weakfloat
     def promote(t):
       if t._uop.base.is_invalid: return t  # invalid bool is weak const
-      dt = weak_dtype(out_dtype) if (weak:=t.dtype in dtypes.weaks) else out_dtype
-      # constuct new const directly
-      return t._wrap_uop(t._uop.const_like(t._uop.base.val, dt)) if weak and t._uop.base.op is Ops.CONST else t.cast(dt)
+      if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const_like(t._uop.base.val, weak_dtype(out_dtype)))
+      return t.cast(out_dtype)
     return promote(x), promote(y)
 
   def _binop(self, op: Ops, x: Self | ConstType, reverse: bool) -> Self:
@@ -398,8 +397,9 @@ class ElementwiseMixin(CreationMixin):
     """
     t, x = self._broadcasted(x)
     # NOTE: the int inverse is done in python, since const has weak dtype without width
-    if dtypes.is_float(dt:=least_upper_dtype(t.dtype, x.dtype)): return -(-t).maximum(-x)
-    return (t ^ (k:=dt.const(dt.min+dt.max))).maximum(x ^ k) ^ k
+    # TODO: clean this up once _broadcasted does not promote dtype
+    if dtypes.is_float(dt:=least_upper_dtype(t.dtype, x.dtype)): return -(-t).alu(Ops.MAX, -x)
+    return (t ^ (k:=dt.const(dt.min+dt.max))).alu(Ops.MAX, x ^ k) ^ k
 
   def copysign(self, other: Self | ConstType) -> Self:
     """

tinygrad/uop/render.py

@@ -102,8 +102,12 @@ pm_pyrender_extra = PatternMatcher([
   # explicit trunc ops: `//` and `%` parse as FLOORDIV/FLOORMOD, so render CDIV/CMOD via .alu()
   (UPat(Ops.CDIV, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CDIV, {ctx[x.src[1]]})"),
   (UPat(Ops.CMOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CMOD, {ctx[x.src[1]]})"),
+  # `.where` re-promotes its operands, so render WHERE via .alu() too
+  (UPat(Ops.WHERE, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.WHERE, {ctx[x.src[1]]}, {ctx[x.src[2]]})"),
+  # the binary operators re-promote their operands (a weak src meeting a strong one gets a cast), render those via .alu() too
   (UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, name="x"), lambda ctx,x:
-    strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
+    strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")
+    if x.src[0]._broadcasted(x.src[1]) == x.src else f"{ctx[x.src[0]]}.alu({x.op}, {ctx[x.src[1]]})"),
   (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
   (UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
     ([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),