fix openpilot speed regeression (#12039)

* set local_size=None if special.arg[0]=='i'

* add cast back
This commit is contained in:
Sieds Lykles
2025-09-06 00:05:45 +02:00
committed by GitHub
parent e00858a2c3
commit 8435d2d23b
+2 -2
View File
@@ -82,9 +82,9 @@ class ProgramSpec:
if u.op is Ops.LOAD: self.ins.extend([x.arg for x in u.src[0].toposort() if x.op is Ops.DEFINE_GLOBAL])
if u.op is Ops.SPECIAL:
# NOTE: you have to set local_size and global_size to the base [1,1,1] outside this
if u.arg[0] == 'i': self.local_size = None
special_size = self.local_size if u.arg[0] == 'l' else self.global_size
assert special_size is not None, f"special_size is None but found SPECIAL in uops {u}"
special_size[int(u.arg[-1])] = cast(int, u.src[0].ssimplify()) # TODO: the type here should be sint
if special_size is not None: special_size[int(u.arg[-1])] = cast(int, u.src[0].ssimplify())
self.vars = sorted(self.vars, key=lambda v: v.arg)
self.outs = sorted(dedup(self.outs))
self.ins = sorted(dedup(self.ins))