From 98de58260b32096d53e2612058d24905c45af71f Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Sat, 19 Oct 2024 19:39:25 -0400
Subject: [PATCH] simplify valid itself (#7112)

---
 test/unit/test_simplify_valid_idx.py |  5 ++---
 tinygrad/codegen/uopgraph.py         | 12 ++++++++++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/test/unit/test_simplify_valid_idx.py b/test/unit/test_simplify_valid_idx.py
index b4d2c931f6f98..9fdc94e23bbf5 100644
--- a/test/unit/test_simplify_valid_idx.py
+++ b/test/unit/test_simplify_valid_idx.py
@@ -65,6 +65,7 @@ def check(self, val0, sidx, svalid):
     self.assertEqual(idx.render(simplify=False), sidx)
     self.assertEqual(valid.render(simplify=False), svalid)
 
+  @unittest.skip("need a different way to test conv2d backward")
   def test_conv_backward(self):
     # DEBUG=4 python3 test/test_ops.py TestOps.test_simple_conv2d
     gidx0 = Special("gidx0", 3)
@@ -134,11 +135,9 @@ def test_simplify_within_valid(self):
     valid = (ridx0*3+ridx1).lt(8) & (((ridx0*3+ridx1)//8+ridx2*3+ridx3)%4).lt(2)
     idx = ridx0+ridx1+ridx2+ridx3
     load = get_gated_load_uop(valid, idx)
-    # TODO: simplify the valid
-    # alu0 = ((ridx0*3)+ridx1)
     self.check(load,
       "(((ridx0+ridx1)+ridx2)+ridx3)",
-      "((((ridx0*3)+ridx1)<8)&(((((((ridx0*3)+ridx1)//8)+(ridx2*3))+ridx3)%4)<2))")
+      "((((ridx0*3)+ridx1)<8)&((((ridx2*3)+ridx3)%4)<2))")
 
 class TestImageSimplification(unittest.TestCase):
   def test_idx_gt_c(self):
diff --git a/tinygrad/codegen/uopgraph.py b/tinygrad/codegen/uopgraph.py
index ec84df4935194..f8b6a9ccb3509 100644
--- a/tinygrad/codegen/uopgraph.py
+++ b/tinygrad/codegen/uopgraph.py
@@ -132,6 +132,16 @@ def uop_given_valid(valid:UOp, uop:UOp) -> Optional[UOp]:
 
   return uop
 
+def simplify_valid(valid:UOp) -> Optional[UOp]:
+  ret:List[UOp] = []
+  something_changed = False
+  try:
+    for stmt in split_uop(valid, BinaryOps.AND):
+      ret.append(stmt if not ret else uop_given_valid(functools.reduce(operator.and_, ret), stmt))
+      if ret[-1] is not stmt: something_changed = True
+    return functools.reduce(operator.and_, ret) if something_changed else None
+  except ValueError: return None
+
 def simplify_buffer_load(load:UOp) -> Optional[UOp]:
   if not isinstance(load.src[0].dtype, PtrDType) or len(load.src) != 4: return None
   buf, start_idx, invalid_val, valid = load.src
@@ -539,6 +549,8 @@ def find_gate(x:UOp) -> Optional[UOp]:
   (UPat(UOps.STORE, name="root"), delete_redundant_gates),
   # late fixup of unfoldable image loads
   (UPat(UOps.LOAD, src=(UPat.var("buf"), UPat()), allow_any_len=True, name="load"), fix_unfoldable_image_load),
+  # simplify valid
+  (UPat(UOps.ALU, name="valid", arg=BinaryOps.AND), simplify_valid),
   # image load valid idx simplification
   (UPat(UOps.LOAD, name="load"), simplify_image_load),
   # buffer load valid idx simplification