update SPECIAL spec (#16952)

also updated tests with bad UOp
2026-07-09 19:05:34 -04:00
parent 95681f17ee
commit 78223d690a
4 changed files with 21 additions and 18 deletions
@@ -34,7 +34,7 @@ class TestRendererFailures(unittest.TestCase):
  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu(self):
    a = UOp.param(0, dtypes.int, (4,))
-    gate_alu = (lidx0:=UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 4),), arg='lidx0')).ne(0)
+    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
    ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
@@ -43,8 +43,8 @@ class TestRendererFailures(unittest.TestCase):
  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu_2d(self):
    a = UOp.param(0, dtypes.int, (8,))
-    gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 4),), arg='lidx0')).ne(0)
-    gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 2),), arg='lidx1')).ne(0)
+    gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
+    gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0)
    gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
    ret = _test_uop_result([], sink, local_size=[4, 2, 1])[0]
@@ -88,7 +88,7 @@ class TestWGSLFailures(unittest.TestCase):
    a = UOp.param(0, dtypes.int, (4,))
    b = UOp.param(1, dtypes.int, (4,))
    c = UOp.param(2, dtypes.int, (4,))
-    lidx0 = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 4),), arg="lidx0")
+    lidx0 = UOp.special(4, "lidx0")
    gate = lidx0.ne(0)
    alt = c.index(lidx0).load()
    ld = UOp.load(b.index(lidx0.valid(gate)))
@@ -103,7 +103,7 @@ class TestPTXFailures(unittest.TestCase):
  @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
  def test_gated_store_with_if(self):
    a = UOp.param(0, dtypes.int, (4,))
-    gate_alu = (lidx0:=UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 4),), arg='lidx0')).ne(0)
+    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    val = UOp.const(dtypes.int, 1)
    if_uop = UOp(Ops.IF, src=(gate_alu,))
    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0, if_uop), val))
@@ -532,7 +532,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_fold_gated_load_local(self):
    glbl0 = UOp.param(0, dtypes.int, (16,))
    smem = UOp.placeholder((18,), dtypes.int, slot=0, addrspace=AddrSpace.LOCAL)
-    lidx = UOp.special(16, "lidx0", dtypes.int)
+    lidx = UOp.special(16, "lidx0")
    st = smem.index(lidx).store(glbl0.index(lidx).load())
    barrier = st.barrier()
    ld0 = smem.after(barrier).index(UOp.invalid())
@@ -111,9 +111,9 @@ class TestExecALU(unittest.TestCase):
 class TestGatedStoreRewrite(unittest.TestCase):
  def test_tiny_gate_store(self):
    gmem = UOp.param(0, dtypes.float, (8,))
-    gidx0 = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 4),), arg='gidx0')
-    gate = gidx0<UOp.const(dtypes.int, 1)
-    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(dtypes.int, 2)).valid(gate)))
+    gidx0 = UOp.special(4, 'gidx0')
+    gate = gidx0<UOp.const(dtypes.index, 1)
+    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(dtypes.index, 2)).valid(gate)))
    val = UOp.const(dtypes.float, 42.0)
    store = UOp(Ops.STORE, src=(idx, val))
    uops = to_uops_list([store])
@@ -128,9 +128,9 @@ class TestGatedStoreRewrite(unittest.TestCase):
  def test_gate_some_stores(self):
    gmem0 = UOp.param(0, dtypes.float, (8,))
    gmem1 = UOp.param(1, dtypes.float, (8,))
-    gidx0 = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 4),), arg='gidx0')
-    idx = gidx0 * UOp.const(dtypes.int, 2)
-    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(dtypes.int, 1))))
+    gidx0 = UOp.special(4, 'gidx0')
+    idx = gidx0 * UOp.const(dtypes.index, 2)
+    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(dtypes.index, 1))))
    idx1 = UOp(Ops.INDEX, src=(gmem1, idx))
    val = UOp.const(dtypes.float, 42.0)
    stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
@@ -148,9 +148,9 @@ class TestGatedStoreRewrite(unittest.TestCase):
  def test_merge_ifs_alt(self):
    gmem0 = UOp.param(0, dtypes.float, (8,))
    gmem1 = UOp.param(1, dtypes.float, (8,))
-    gidx0 = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 4),), arg='gidx0')
-    idx = gidx0*UOp.const(dtypes.int, 2)
-    gate = gidx0<UOp.const(dtypes.int, 1)
+    gidx0 = UOp.special(4, 'gidx0')
+    idx = gidx0*UOp.const(dtypes.index, 2)
+    gate = gidx0<UOp.const(dtypes.index, 1)
    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gate)))
    idx1 = UOp(Ops.INDEX, src=(gmem1, idx.valid(gate)))
    val = UOp.const(dtypes.float, 42.0)
@@ -103,9 +103,6 @@ spec_shared = PatternMatcher([
  # BARRIER (on any length). TODO: this should only be in spec_program
  (UPat(Ops.BARRIER, dtypes.void), lambda: True),

-  # SPECIAL. TODO: this should only be in spec_program
-  (UPat(Ops.SPECIAL, src=(UPat.var("x", (dtypes.index, dtypes.int32)),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
-
  # assembly instruction
  (UPat(Ops.INS), lambda: True),

@@ -150,6 +147,9 @@ spec_tensor = PatternMatcher([
  (UPat(Ops.GETTUPLE, src=(UPat(Ops.FUNCTION, src=(UPat(Ops.TUPLE, name="t"),), allow_any_len=True),), name="g"), valid_gettuple),
  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), valid_gettuple),

+  # SPECIAL is index before index lowering. custom_kernel currently has this
+  (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.index),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
+
  # inputs to movement ops
  (UPat({Ops.ADD, Ops.MUL, Ops.CDIV, Ops.FLOORDIV}, dtype=dtypes.index), lambda: True),

@@ -211,6 +211,9 @@ spec_program = PatternMatcher([
  # if has a <gate, index_for_dedup>
  (UPat(Ops.IF, dtype=dtypes.void, src=(UPat(dtype=dtypes.bool), UPat((Ops.CAST, Ops.INDEX, Ops.SHRINK)))), lambda: True),
  (UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),)), lambda: True),
+
+  # SPECIAL is int32 after index lowering
+  (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.int32),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
 ])+spec_shared

 # these are intermediate ops. everything should be deleted from here