Merge remote-tracking branch 'origin/master' into qwen36_27b_amd_900

test/backend/test_wait_loop.py

@@ -81,6 +81,20 @@ def loop_in_loop_kernel(C:UOp) -> UOp:
 
   return C[0].store(i[0].load()).sink(arg=KernelInfo(name="loop_in_loop", opts_to_apply=()))
 
+def scalar_alu_index_kernel(out: UOp) -> UOp:
+  val = UOp.param(1, dtypes.int, (1,), vmin_vmax=(0, 100), addrspace=AddrSpace.ALU, name="val")
+  idx = UOp(Ops.INDEX, dtypes.int, (val, UOp.const(0, dtypes.int)))
+  return out[0].store(idx + 1).sink(arg=KernelInfo(name="scalar_alu_index"))
+
+class TestScalarALUIndex(unittest.TestCase):
+  def test_scalar_alu_index(self):
+    # regression: indexing a scalar (1,)-shaped ALU param must render the variable
+    # directly, not as data.x (which is a compile error on a scalar int)
+    out = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize()
+    result = Tensor.custom_kernel(out, fxn=scalar_alu_index_kernel)[0]
+    run_linear(result.schedule_linear(), var_vals={"val": 42})
+    self.assertEqual(result.item(), 43)
+
 class TestWaitLoop(unittest.TestCase):
   def test_wait_loop(self):
     c = Tensor.empty(1, dtype=dtypes.int)

tinygrad/codegen/__init__.py

@@ -150,6 +150,9 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
   (UPat(GroupOp.Elementwise|{Ops.LOAD,Ops.STORE}, name="b"), do_devectorize),
   # INDEX without src is nothing (TODO: this should be in mop_cleanup)
   (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
+  # INDEX of a scalar ALU param is the param itself
+  (UPat(Ops.INDEX, src=(UPat(Ops.PARAM, name="v"),), allow_any_len=True),
+   lambda v: v if v.addrspace == AddrSpace.ALU and v.max_numel() == 1 else None),
   # unpack WMMA
   (UPat(Ops.WMMA, name="u"), do_stack_wmma),
   # stacked INDEX is many INDEX

tinygrad/codegen/decomp/op.py

@@ -47,17 +47,17 @@ def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None:
 
 def threefry2x32(x: UOp, key: UOp):
   # split x and key from uint64 to two uint32
-  x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32)
-  key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32)
+  x0, x1 = x.cast(dtypes.uint32), (x >> 32).cast(dtypes.uint32)
+  key0, key1 = key.cast(dtypes.uint32), (key >> 32).cast(dtypes.uint32)
 
   rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
   ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0]
   xr:list[UOp] = [x0 + ks[-1], x1 + ks[0]]
   for i in range(5):
-    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r)))
+    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] << r) + (xr[1] >> (32 - r)))
     xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)]
 
-  return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64)
+  return (xr[1].cast(dtypes.uint64) << 32) | xr[0].cast(dtypes.uint64)
 
 # ***** decomposition patterns *****
 

tinygrad/mixin/rand.py

@@ -12,7 +12,7 @@ class RandMixin(OpMixin):
   def _threefry_random_bits(key, counts0, counts1):
     x = (counts1.cast(dtypes.uint64) << 32) | counts0.cast(dtypes.uint64)
     x = x.threefry((key[1].cast(dtypes.uint64) << 32) | key[0].cast(dtypes.uint64))
-    return (x & 0xffffffff).cast(dtypes.uint32).cat(((x >> 32) & 0xffffffff).cast(dtypes.uint32))
+    return x.cast(dtypes.uint32).cat((x >> 32).cast(dtypes.uint32))
 
   @classmethod
   def random_bits(cls, key:Self, counter:Self, num:int) -> Self:

tinygrad/runtime/support/usb.py

@@ -90,10 +90,11 @@ class CustomASM24Controller:
   def __init__(self, usb:USB3):
     self.usb = usb
 
-    # Custom firmware now boots with PCIe off. Power it on before probing the link.
-    ltssm = self.read(0xB450, 1)[0]
-    if ltssm != 0x78: self.set_pcie_power(True)
-    ltssm = self.read(0xB450, 1)[0]
+    # Custom firmware now boots with PCIe off. Power it on before probing the link with a 5s grace period.
+    if (ltssm:=self.read(0xB450, 1)[0]) != 0x78:
+      self.set_pcie_power(True)
+      grace_period = time.monotonic() + 5.
+      while time.monotonic() < grace_period and (ltssm:=self.read(0xB450, 1)[0]) != 0x78: time.sleep(0.1)
     if ltssm != 0x78: raise RuntimeError(f"PCIe link not up (LTSSM=0x{ltssm:02X}), custom firmware not ready")
 
   def set_pcie_power(self, enabled:bool, timeout:int=10000): self.usb.control_write(0xF3, value=int(enabled), timeout=timeout)

tinygrad/uop/symbolic.py

@@ -127,7 +127,6 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
   (UPat.var("x") ^ UPat.var("x"), lambda x: x.const_like(0)), # x^x -> 0
   (UPat.var("x") & 0, lambda x: x.const_like(0)), # x&0 -> 0
   # (x&mask)>>k -> x>>k when mask only clears bits below k
-  # TODO: combine this with "# rules for threefry" below
   ((UPat.var("x") & UPat.cvar("mask")) >> UPat.cvar("k"),
    lambda x,mask,k: x >> k.val if mask.val | ((1 << k.val) - 1) == -1 else None),
   ((UPat.var("x") & UPat.cvar("mask")) // UPat.cvar("c"),
@@ -168,13 +167,10 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
   (UPat.var("x").alu(Ops.POW, UPat.cvar("c")), simplify_pow),
   # positive const ** x
   (UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.val == 1 else (x*math.log2(c.val)).exp2() if c.val > 0 else None),
-  # rules for threefry
-  ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)),
-  (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
-  (((UPat.var('x',  dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x),
-  (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
-  (((UPat.var('x',  dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x),
-  (((UPat.var('x',  dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32, lambda x: x),
+  # unpack a uint64 packed from two uint32 (threefry)
+  (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
+  (((UPat.var('x', dtypes.uint32).cast(dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32,
+   lambda x: x.cast(dtypes.uint64)),
   # ** simple where folding **
   # a conditional with the same results either way is a noop, also fold const conditionals
   (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val),