From e3dc10f8f6980259be5ae4ec7ee2a22fed37c9f1 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 10 Oct 2024 10:52:05 -0400 Subject: [PATCH] improve fold_unrolled_divs (#6977) addressed #6935 the first few terms in fold_unrolled_divs might have been folded already, so the check should first try to add those terms back. there is a case that every but one term is folded which is not an add chain anymore, so just added as a failed test case for now --- test/unit/test_uop_symbolic.py | 15 +++++++++++++++ tinygrad/ops.py | 12 +++++++++--- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/test/unit/test_uop_symbolic.py b/test/unit/test_uop_symbolic.py index 173bda904bb21..aba8aa5240bd8 100644 --- a/test/unit/test_uop_symbolic.py +++ b/test/unit/test_uop_symbolic.py @@ -438,6 +438,21 @@ def test_arange_unrolled4(self): unrolled_div = (gidx+2561)//4+(gidx+2562)//4+(gidx+2560)//4+(gidx+2559)//4 self.helper_test_variable(unrolled_div, 2559, 5118, "(gidx+2559)") + def test_arange_unrolled4_small(self): + gidx = Variable("gidx", 0, 3) + unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4 + self.helper_test_variable(unrolled_div, 0, 3, "gidx") + + gidx = Variable("gidx", 0, 2) + unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4 + self.helper_test_variable(unrolled_div, 0, 2, "gidx") + + # TODO: fix this, it has only one term and is no longer an add chain + with self.assertRaises(AssertionError): + gidx = Variable("gidx", 0, 1) + unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4 + self.helper_test_variable(unrolled_div, 0, 1, "gidx") + def test_arange_unrolled2(self): gidx = Variable("gidx", 0, 2559) unrolled_div = (gidx+2559)//2+(gidx+2560)//2+3 diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 12ef986968c03..f3bed7d1ab43d 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -823,9 +823,11 @@ def lt_folding(x:UOp, c:int) -> Optional[UOp]: def fold_unrolled_divs(divs:UOp): # div pattern in unrolled arange # example: (x//4+(x+1)//4+(x+2)//4+(x+3)//4 -> x - add_chain, seen_const, ans = list(_get_chain(divs, BinaryOps.ADD)), [], None + add_chain, denominator, seen_const, ans = list(_get_chain(divs, BinaryOps.ADD)), None, [], None for u in add_chain: - if not (u.op is UOps.ALU and u.arg is BinaryOps.IDIV and u.src[1].op is UOps.CONST and u.src[1].arg==len(add_chain)): return None + if not (u.op is UOps.ALU and u.arg is BinaryOps.IDIV and u.src[1].op is UOps.CONST): return None + if denominator is None: denominator = u.src[1].arg + if denominator != u.src[1].arg: return None # assumed CONST is the last of an ADD if (s0:=u.src[0]).op is UOps.ALU and s0.arg is BinaryOps.ADD and s0.src[1].op is UOps.CONST and s0.src[1].op is UOps.CONST: seen_const.append(s0.src[1].arg) @@ -833,7 +835,11 @@ def fold_unrolled_divs(divs:UOp): else: seen_const.append(0) if ans is None: ans = s0 if ans is not s0: return None - return ans if ans is not None and sorted(seen_const)==list(range(len(add_chain))) else None + if denominator is None: return None + # the first (denominator-len(seen_const)) terms may have been folded to 0 already + for i in range(denominator-len(seen_const)): + if ans is not None and 0 <= ans.vmin and ans.vmax + i < denominator: seen_const.append(i) + return ans if ans is not None and sorted(seen_const)==list(range(denominator)) else None def is_irreducible(u:UOp): return u.op in (UOps.DEFINE_VAR, UOps.SPECIAL, UOps.RANGE)