render casted const [pr]

test/backend/test_encodings.py

@@ -2,7 +2,7 @@
 from tinygrad import Device
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.dtype import dtypes
-from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg
+from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, machine_const, def_reg
 
 def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag)
 
@@ -14,51 +14,51 @@ def encode(self, u:UOp): return Device[Device.DEFAULT].renderer.render([u])
 
   # displacement of 0 isn't emitted
   def test_base_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RDI), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RDI)
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RDI), UOp(Ops.NOOP), machine_const(0), machine_const(4)), RDI)
     # mov edi, dword ptr [rdi]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 3F"))
 
   # rsp/r12 require a sib byte when used as base memory address
   def test_rsp_base_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RSP)
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), machine_const(0), machine_const(4)), RSP)
     # mov esp, dword ptr [rsp]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 24 24"))
 
   # rbp/r13 require a displacement when used as base memory address
   def test_rbp_base_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RBP), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RBP)
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RBP), UOp(Ops.NOOP), machine_const(0), machine_const(4)), RBP)
     # mov ebp, dword ptr [rbp + 0]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 6D 00"))
 
   # test [base + index*scale]
   def test_base_index_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, RDX), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RAX)
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, RDX), machine_const(0), machine_const(4)), RAX)
     # mov eax, dword ptr [rax + rdx*4]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 04 90"))
 
   # rsp as index means no index
   def test_rsp_index_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, RSP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RAX)
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, RSP), machine_const(0), machine_const(4)), RAX)
     # mov eax, dword ptr [rax]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 00"))
 
   # however r12 is a valid index
   def test_r12_index_address(self):
     load = ins(X86Ops.MOV, dtypes.int32,
-               (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, GPR[12]), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RAX)
+               (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, GPR[12]), machine_const(0), machine_const(4)), RAX)
     # mov eax, dword ptr [rax + r12*4]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("42 8B 04 A0"))
 
   # test [base + index*scale + 8bit disp]
   def test_complex_address_8bit_disp(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4)), RDI)
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), machine_const(10), machine_const(4)), RDI)
     # mov edi, dword ptr [rdi + rsi*4 + 0xa]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 7C B7 0A"))
 
   # test [base + index*scale + 32bit disp]
   def test_complex_address_32bit_disp(self):
     load = ins(X86Ops.MOV, dtypes.int32,
-               (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10000), imm(dtypes.uint8, 4)), RDI)
+               (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), machine_const(10000), machine_const(4)), RDI)
     # mov edi, dword ptr [rdi + rsi*4 + 0x2710]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B BC B7 10 27 00 00"))
 
@@ -116,30 +116,57 @@ def test_reg_in_imm_field(self):
 
   # when writting to mem the uop takes the store form where dtype is void and there's no definition
   def test_write_mem(self):
-    address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4))
+    address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), machine_const(10), machine_const(4))
     xmm0 = def_reg(dtypes.float32, XMM[0])
-    extr = ins(X86Ops.VPEXTRD, dtypes.void, address + (xmm0, imm(dtypes.uint8, 0)))
+    extr = ins(X86Ops.VPEXTRD, dtypes.void, address + (xmm0, machine_const(0)))
     # vpextrd dword ptr [rdi + rsi*4 + 0xa], xmm0, 0
     self.assertEqual(bytes.fromhex(self.encode(extr)), bytes.fromhex("C4 E3 79 16 44 B7 0A 00"))
 
   # test two address instruction with fused load works
   def test_two_address_load(self):
-    address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4))
+    address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), machine_const(10), machine_const(4))
     cmove = ins(X86Ops.CMOVE, dtypes.int32, address, RAX)
     # cmove eax, dword ptr [rdi + rsi*4 + 0xa]
     self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 44 B7 0A"))
 
   # test instruction where displacement and imm have the same value
   def test_disp_imm_same_value(self):
-    address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int8, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 1))
-    mov = ins(X86Ops.MOVi, dtypes.void, address + (imm(dtypes.int8, 10),))
+    address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int8, RSI), machine_const(10), machine_const(1))
+    mov = ins(X86Ops.MOVi, dtypes.void, address + (machine_const(10),))
     # mov byte ptr [rdi + rsi + 0xa], 0xa
     self.assertEqual(bytes.fromhex(self.encode(mov)), bytes.fromhex("40 C6 44 37 0A 0A"))
 
-    address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10), imm(dtypes.uint8, 4))
-    imul = ins(X86Ops.IMULi, dtypes.int32, address + (imm(dtypes.int32, 10),), RDI)
+    address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), machine_const(10), machine_const(4))
+    imul = ins(X86Ops.IMULi, dtypes.int32, address + (machine_const(10),), RDI)
     # imul edi, dword ptr [rdi + rsi*4 + 0xa], 0xa
-    self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 BC B7 0A 00 00 00 0A 00 00 00"))
+    self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 7C B7 0A 0A 00 00 00"))
+
+  # a displacement takes the smallest width that holds it wherever it was minted, so a stack argument offset is a disp8 like any other
+  def test_disp_width(self):
+    def load(d): return ins(X86Ops.MOV, dtypes.uint64, (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), machine_const(d), machine_const(8)), RAX)
+    # mov rax, qword ptr [rsp + 8] / [rsp + 127]
+    self.assertEqual(bytes.fromhex(self.encode(load(8))), bytes.fromhex("48 8B 44 24 08"))
+    self.assertEqual(bytes.fromhex(self.encode(load(127))), bytes.fromhex("48 8B 44 24 7F"))
+    # the rule is on the magnitude, so 128 and -128 both take 4 bytes
+    self.assertEqual(bytes.fromhex(self.encode(load(128))), bytes.fromhex("48 8B 84 24 80 00 00 00"))
+    self.assertEqual(bytes.fromhex(self.encode(load(-128))), bytes.fromhex("48 8B 84 24 80 FF FF FF"))
+
+  # a machine const is a number with no width, the field it lands in writes its low bytes
+  def test_imm_low_bytes(self):
+    # movabs rax, 0xfffffffffffffffe, the value is wider than the register
+    movabs = ins(X86Ops.MOVABS, dtypes.uint64, (machine_const(2**65-2),), RAX)
+    self.assertEqual(bytes.fromhex(self.encode(movabs)), bytes.fromhex("48 B8 FE FF FF FF FF FF FF FF"))
+    for dt,v,b in ((dtypes.int8, 300, "2C"), (dtypes.int8, -200, "38"), (dtypes.uint8, 300, "2C")):
+      # mov al, <low byte>
+      self.assertEqual(bytes.fromhex(self.encode(ins(X86Ops.MOVi, dt, (machine_const(v),), RAX))), bytes.fromhex("40 C6 C0 " + b))
+    # mov ax, 0x1170
+    self.assertEqual(bytes.fromhex(self.encode(ins(X86Ops.MOVi, dtypes.int16, (machine_const(70000),), RAX))), bytes.fromhex("66 40 C7 C0 70 11"))
+
+  # a constant in an immediate slot takes the width of the slot, not of the value, in either honest form
+  def test_pair_imm_width(self):
+    # add rax, 5 with an imm32, the operand being 64 bit doesn't widen the immediate
+    for c in (machine_const(5), UOp.const(5).cast(dtypes.int64)):
+      self.assertEqual(bytes.fromhex(self.encode(ins(X86Ops.ADDi, dtypes.int64, (c,), RAX))), bytes.fromhex("48 81 C0 05 00 00 00"))
 
   # cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding
   def test_cmove_ignore_cmp(self):
@@ -148,4 +175,4 @@ def test_cmove_ignore_cmp(self):
     self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))
 
 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()

test/backend/test_isel.py

@@ -1,10 +1,9 @@
 import unittest
 from typing import cast
 from tinygrad import Device
-from tinygrad.uop import Ops
 from tinygrad.uop.ops import UOp, dtypes, graph_rewrite
-from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
-from tinygrad.renderer.isa import IselContext
+from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops, to_imm
+from tinygrad.renderer.isa import IselContext, machine_const
 
 # INDEX on a register value with a constant index extracts a single element (the old GEP)
 def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype.scalar())
@@ -48,8 +47,18 @@ def test_complex_address(self):
     a = UOp.variable("a", 0, 0, dtypes.int32)
     load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).load()
     n = self.isel_rewrite(load)
-    # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
-    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 4)
+    # displacement is the constant in "a" scaled to the buffer element size; encoding chooses its width
+    self.assertIs(n.src[2], machine_const(4))
+
+  # a constant reaches isel as a strong CONST or as a cast weak one, both select the same machine const
+  def test_to_imm_forms(self):
+    for dt in (dtypes.int8, dtypes.uint8, dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64):
+      for v in (0, 1, dt.min, dt.max): self.assertIs(to_imm(UOp.const(v, dt)), to_imm(UOp.const(v).cast(dt)), f"{dt} {v}")
+
+  # a machine const states the number the hardware sees, and only fits an immediate if it fits 4 bytes
+  def test_to_imm_value(self):
+    self.assertIs(to_imm(UOp.const(-1, dtypes.uint32)), machine_const(2**32-1))
+    for dt,v in ((dtypes.int64, 2**31), (dtypes.uint64, 2**32), (dtypes.weakint, 2**31)): self.assertIsNone(to_imm(UOp.const(v, dt)))
 
 if __name__ == "__main__":
   unittest.main()

test/mockgpu/amd/pcode.py

@@ -7,6 +7,10 @@
 # Type alias for vars dict: stores UOps and tuples for lambda definitions
 VarVal = UOp | tuple[str, list[str], str]
 
+# a committed constant reaches the emulator in either honest form: the typed CONST or a CAST over one.
+# the inner CONST is bare for a tinygrad pair and typed for the emulator's own IR, so neither dtype is required
+def _cc(u): return u.op is Ops.CONST or (u.op is Ops.CAST and u.src[0].op is Ops.CONST)
+def _cv(u): return u.src[0].val if u.op is Ops.CAST else u.val
 def _const(dt, v): return UOp.const(v, dt)
 def _u32(v): return _const(dtypes.uint32, v)
 def _u64(v): return _const(dtypes.uint64, v)
@@ -55,8 +59,8 @@ def _expr_bits(v: UOp) -> int:
   if v.op in (Ops.AND, Ops.XOR):
     widths: list[int] = []
     for src in v.src:
-      if src.op == Ops.CONST and isinstance(src.val, int) and src.val > 0 and (src.val & (src.val + 1)) == 0:
-        widths.append(src.val.bit_length())
+      if _cc(src) and isinstance(_cv(src), int) and _cv(src) > 0 and (_cv(src) & (_cv(src) + 1)) == 0:
+        widths.append(_cv(src).bit_length())
     if widths: return max(widths)
   return v.dtype.bitsize
 
@@ -144,9 +148,9 @@ def minmax(a: UOp, b: UOp) -> UOp:
 def _find_two_pi_mul(x):
   if x.op != Ops.MUL or len(x.src) != 2: return None
   for i, s in enumerate(x.src):
-    if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
+    if _cc(s) and abs(_cv(s) - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
     if s.op == Ops.MUL and len(s.src) == 2:
-      vals = [ss.val for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].val for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
+      vals = [_cv(ss) for ss in s.src if _cc(ss)]
       if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1])
   return None
 
@@ -163,7 +167,7 @@ def _trig_reduce(x, phase=0.0):
 
 def _signext(val: UOp) -> UOp:
   for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
-    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].val == mask) or val.dtype.itemsize == bits // 8:
+    if (val.op == Ops.AND and len(val.src) == 2 and _cc(val.src[1]) and _cv(val.src[1]) == mask) or val.dtype.itemsize == bits // 8:
       v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
       sb = (v32 >> _u32(bits - 1)) & _u32(1)
       return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
@@ -497,7 +501,7 @@ def _apply_binop(self, left, right, op):
         if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32)
         return (left >> right) if op == '>>' else (left << right)
       case '+' | '-':
-        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val)
+        if op == '-' and _cc(left) and _cc(right): return _const(left.dtype, _cv(left) - _cv(right))
         return (left + right) if op == '+' else (left - right)
       case '*' | '/':
         # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow
@@ -507,7 +511,7 @@ def _apply_binop(self, left, right, op):
           left, right = left.cast(pdt), right.cast(pdt)
         if op == '*': return left * right
         return (left // right) if dtypes.is_int(left.dtype) else (left / right)
-      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 2.0 else left
+      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if _cc(left) and _cv(left) == 2.0 else left
 
   _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]
 
@@ -529,8 +533,8 @@ def unary(self) -> UOp:
       return inner.eq(_const(inner.dtype, 0))
     if self.try_eat_val('-', 'OP'):
       inner = self.unary()
-      if inner.op == Ops.CONST:
-        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val)
+      if _cc(inner):
+        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -_cv(inner))
       return inner.neg()
     if self.try_eat_val('+', 'OP'): return self.unary()
     return self.postfix()
@@ -669,15 +673,15 @@ def _handle_bracket_rest(self, first: UOp, base: UOp, var_name: str | None = Non
       self.eat('OP')
       width = self.parse()
       self.eat('RBRACKET')
-      if width.op == Ops.CONST:
-        w = int(width.val)
+      if _cc(width):
+        w = int(_cv(width))
         return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1)
       return base
     if self.try_eat('COLON'):
       second = self.parse()
       self.eat('RBRACKET')
-      if first.op == Ops.CONST and second.op == Ops.CONST:
-        a, b = int(first.val), int(second.val)
+      if _cc(first) and _cc(second):
+        a, b = int(_cv(first)), int(_cv(second))
         if a < b: return _bitreverse(base, b - a + 1)
         hi, lo = a, b
         if lo >= base.dtype.itemsize * 8:
@@ -698,8 +702,8 @@ def _handle_bracket_rest(self, first: UOp, base: UOp, var_name: str | None = Non
       dt_suffix = DTYPES.get(self.eat('IDENT').val, dtypes.uint32)
     if var_name is None:
       var_name = self._find_var_name(base)
-    if first.op == Ops.CONST:
-      idx = int(first.val)
+    if _cc(first):
+      idx = int(_cv(first))
       # Check for array element (var@idx)
       if var_name and f'{var_name}@{idx}' in self.vars:
         v = self.vars[f'{var_name}@{idx}']
@@ -872,7 +876,7 @@ def mindex(idx:UOp): return mem.index(idx.valid(active) if active is not None el
 
   def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
     if l.dtype != r.dtype:
-      if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 0: l = l.cast(dtypes.int)
+      if r.dtype == dtypes.int and _cc(r) and _cv(r) < 0: l = l.cast(dtypes.int)
       else: r = r.cast(l.dtype)
     return l, r
 
@@ -969,8 +973,8 @@ def parse_bound():
           p.eat('QUOTE')
         if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
         expr = p.parse().simplify()
-        assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
-        return int(expr.val)
+        assert _cc(expr), f"loop bound must be constant, got {expr}"
+        return int(_cv(expr))
       start_val = parse_bound()
       p.eat('COLON')
       end_val = parse_bound()
@@ -1258,7 +1262,7 @@ def parse_bound():
       def parse_cond(s, kw):
         ll = s.lower()
         return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), env, funcs))
-      def is_const(c, v): return c.op == Ops.CONST and c.val is v
+      def is_const(c, v): return _cc(c) and _cv(c) is v
       cond = parse_cond(line, 'if')
       conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else []
       branch_assigns: list[tuple[UOp, list]] = []  # (cond, assigns_list) for side-effect merging
@@ -1341,9 +1345,9 @@ def _cond_side_effect(cnd, dest, val):
           # Build combined condition: each branch fires when its cond is true AND no earlier cond was true
           remaining = UOp.const(True)
           for bc, bse in branch_assigns:
-            effective = remaining & bc if remaining.op != Ops.CONST else bc
+            effective = remaining & bc if not _cc(remaining) else bc
             for dest, val in bse: assigns.append(_cond_side_effect(effective, dest, val))
-            remaining = remaining & bc.logical_not() if remaining.op != Ops.CONST else bc.logical_not()
+            remaining = remaining & bc.logical_not() if not _cc(remaining) else bc.logical_not()
           for dest, val in else_side_effects: assigns.append(_cond_side_effect(remaining, dest, val))
       continue
 

test/null/test_gpudims.py

@@ -27,7 +27,7 @@ def _verify_indices_z3(self, idxs, dims):
     flat = UOp.const(0)
     for i, idx in enumerate(idxs):
       flat = flat + idx * int(math.prod(dims[i+1:]))
-    flat_p = flat.substitute({s: UOp(Ops.SPECIAL, src=s.src, arg=s.arg+"_p") for s in specials})
+    flat_p = flat.substitute({s: UOp(Ops.SPECIAL, s.dtype, src=s.src, arg=s.arg+"_p") for s in specials})
     solver = z3.Solver()
     [z3_flat, z3_flat_p] = uops_to_z3(solver, flat, flat_p)
     # bounds