clean up divide_exact order [pr] (#12919)

do the const first since ADD can also call into that
This commit is contained in:
chenyu
2025-10-25 18:47:57 -04:00
committed by GitHub
parent e18922f111
commit 94701d4838
+1 -1
View File
@@ -657,8 +657,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
return math.prod([*count.elements(), terms[0].const_like(math.gcd(*factors))]) # put the const at the top
def divide_exact(self, v:UOp) -> UOp|None:
if self is v: return self.const_like(1)
if self.op is Ops.ADD: return None if (s0:=self.src[0].divide_exact(v)) is None or (s1:=self.src[1].divide_exact(v)) is None else s0+s1
if v.op is Ops.CONST: return self.divides(v.arg)
if self.op is Ops.ADD: return None if (s0:=self.src[0].divide_exact(v)) is None or (s1:=self.src[1].divide_exact(v)) is None else s0+s1
if self.op is Ops.MUL:
(fac, const), (div_fac, div_const) = self.pop_const(Ops.MUL), v.pop_const(Ops.MUL)
new_count = collections.Counter(fac.split_uop(Ops.MUL))