Repository: tinygrad/tinygrad
Pull request: Fix VALIDHACKS for Images and make it default
So this is just a different approach to validhacks.

Fixes #1801

---------
Notes:
* Some cases indexes have sum node but still requires valid. It should not require it as it needs to go out of bound ?

Problem statement:
Fix VALIDHACKS for Images and make it default ($300 bounty)
When you read images out of bounds, they will return 0s. Currently the compiler is unaware of this and still gates the load. Figure out when we don't need it and disable it.

Images are used in the openpilot model `openpilot/go.sh` that have this extra gated load. Safely remove it!

Must be well tested for bounty, it's easy to do this subtly wrong.

Simple example of issue:
`GPU=1 DEBUG=4 FORWARD_ONLY=1 IMAGE=2 python3 test/test_ops.py TestOps.test_simple_padding_conv2d`

generates 

`float4 val0 = ((((lidx0*(-1))<0)*(lidx0<3)))?(read_imagef(data1, smp, (int2)(((lidx0+1)%2),(((lidx0+1)/2)+(-1))))):(float4)(0.0f,0.0f,0.0f,0.0f); # (lidx0 ranges from 0-3)`

instead of

`float4 val0 = read_imagef(data1, smp, (int2)(lidx0-1,0))`

to read image

`dtypes.imagef((1, 2, 4)) # the last 4 is the float4, this is a 2x1 image`

That gate is not needed if you remove the `%2` and subtract 2 from the index. You also then don't need the y index at all.

See `validhacks` in `to_image_idx` for the old (broken) code that hacked this. The symbolic engine should be good enough now to do this properly.

Hints from issue discussion:
If you'd like some help from GPT-4 to understand this:

https://chat.openai.com/share/cabc655a-6cc7-4afa-bda4-741d7600eb2c
I was just experimenting with the test_simple_padding_conv2d test and noticed a max operation that is performed on two equal constants. Since this might be optimized for constants or other equal variables, I think it is useful to know. I couldn't find something in lazy.py to optimize it. The following output is obtained with forward_pass=1.

Edit: fixed this in https://github.com/tinygrad/tinygrad/pull/1858
```
__kernel void E_4_3(__global float* data0, const __global float* data1, const __global float* data2) {
float val0 = data1[0];
float val1 = data1[1];
float val2 = data2[0];
float alu0 = max(0.0f,0.0f);
*((__global float4*)(data0+0)) = (float4)(float4)(alu0,alu0,alu0,alu0);
*((__global float4*)(data0+4)) = (float4)(float4)(alu0,max((val0*val2),0.0f),max((val1*val2),0.0f),alu0);
*((__global float4*)(data0+8)) = (float4)(float4)(alu0,alu0,alu0,alu0);
}
```

Commit message: Managed to reduce openpilot time from 30 secs to 5 secs

Patch to review:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index dd669dacc4bb4..f874ac1b0c469 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -147,7 +147,7 @@ jobs:
       - if: ${{ matrix.task == 'openpilot' }}
         name: Test openpilot model compile and size
         run: |
-          DEBUG=2 ALLOWED_KERNEL_COUNT=209 FLOAT16=1 DEBUGCL=1 GPU=1 IMAGE=2 python openpilot/compile.py
+          DEBUG=2 ALLOWED_KERNEL_COUNT=209 VALIDTEST=1 FLOAT16=1 DEBUGCL=1 GPU=1 IMAGE=2 python openpilot/compile.py
           python -c 'import os; assert os.path.getsize("/tmp/output.thneed") < 100_000_000'
       - if: ${{ matrix.task == 'openpilot' }}
         name: Test openpilot model correctness (float32)
diff --git a/openpilot/compile.py b/openpilot/compile.py
index 17bada5da542d..708fe6fe8bb41 100644
--- a/openpilot/compile.py
+++ b/openpilot/compile.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-import os, time, io, pathlib, sys, traceback
+import os, time, io, pathlib, sys, traceback, re
 sys.path.insert(0, str(pathlib.Path(__file__).parents[1]))
 
 if os.getenv("OPT", None) is None:
@@ -72,6 +72,11 @@ def compile(dat, output_fn):
     # pass these to thneed
     setattr(prg.clprg, 'op_estimate', prg.op_estimate)
     setattr(prg.clprg, 'prg', prg.prg)
+
+    if getenv("VALIDTEST") == 1:
+      src = re.search(r"=.*\?.*?read_image", prg.prg)
+      if src is not None: raise Exception("Openpilot has valid checks!")
+
     global_size = prg.global_size + [1]*(3-len(prg.global_size))
     local_size = prg.local_size + [1]*(3-len(prg.local_size))
     cl_cache.append((prg.clprg, [[g*l for g,l in zip(global_size, local_size)], local_size, *[x._buf for x in args]]))
diff --git a/test/unit/test_symbolic.py b/test/unit/test_symbolic.py
index 57fff2250124a..64ff1b590f4f0 100644
--- a/test/unit/test_symbolic.py
+++ b/test/unit/test_symbolic.py
@@ -130,6 +130,10 @@ def test_mod_factor(self):
     # NOTE: even though the mod max is 50, it can't know this without knowing about the mul
     self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*100, Variable("b", 0, 3)*50]) % 100, 0, 99, "((b*50)%100)")
 
+  def test_mod_to_sub(self):
+    # This is mod reduction
+    self.helper_test_variable((1+Variable("a",1,2))%2, 0, 1, (Variable("a",1,2)-1).render())
+
   def test_sum_div_const(self):
     self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*4, Variable.num(3)]) // 4, 0, 7, "a")
 
diff --git a/tinygrad/codegen/linearizer.py b/tinygrad/codegen/linearizer.py
index eb0a45e1c1899..7704b924e2879 100644
--- a/tinygrad/codegen/linearizer.py
+++ b/tinygrad/codegen/linearizer.py
@@ -4,12 +4,12 @@
 from collections import defaultdict
 from enum import Enum, auto
 
-from tinygrad.helpers import colored, ImageDType, DEBUG, dtypes, DType, partition, prod, PtrDType, all_same
+from tinygrad.helpers import colored, ImageDType, DEBUG, dtypes, DType, prod, PtrDType, all_same, partition
 from tinygrad.ops import LazyOp, UnaryOps
 from tinygrad.ops import ReduceOps, BinaryOps, TernaryOps
 from tinygrad.runtime.lib import RawConst
 from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.symbolic import Variable, NumNode, Node, SumNode, MulNode, DivNode, ModNode, LtNode, AndNode, sym_rename
+from tinygrad.shape.symbolic import Variable, NumNode, Node, SumNode, MulNode, DivNode, ModNode, LtNode, AndNode, sym_rename, sym_render
 from tinygrad.codegen.optimizer import OptimizedKernel
 from tinygrad.codegen.kernel import LocalBuffer
 VariableOrNum = Union[Variable, NumNode, Node]
@@ -21,25 +21,158 @@ class UOps(Enum):
   LOAD = auto(); STORE = auto(); CONST = auto(); BARRIER = auto() # noqa: E702
   ALU = auto(); WMMA = auto(); CAST = auto(); GEP = auto() # noqa: E702
 
-def to_image_idx(base_shape:Tuple[int, ...], idxy:Node, valid:Node, validhacks=False) -> Tuple[Node, Node]:
-  idy = (idxy//(4*base_shape[1]))
-  if validhacks and valid.min == 0:
-    idx = (idxy//4) + (idy*-base_shape[1])
-    # find the ones in idx that didn't factorize and remove them (TODO: this is not universal)
-    if isinstance(idx, SumNode):
-      unfactored, idx_nodes = partition(idx.nodes, lambda x: isinstance(x, MulNode) and x.b == -base_shape[1])
-      assert len(unfactored) <= 1
-      idx = Variable.sum(idx_nodes)
-      unfactored = (Variable.sum(unfactored) // base_shape[1])
-      idy += unfactored
-      # ugh really...handtuned garbage
-      if idx.min >= (base_shape[1]*3)//4:
-        idx -= base_shape[1]
-        idy += 1
-  else:
-    idx = (idxy//4)%base_shape[1]
-  if DEBUG >= 5: print("to_image_idx", base_shape, idx.min, idx.max, idy.min, idy.max, idx, idy)
-  return idx, idy
+def to_image_idx(base_shape:Tuple[int, ...], idxy:Node, valid:Node) -> Tuple[Tuple[Node, Node], Node]:
+  if valid.min == 0:
+    nodes: List = valid.nodes if isinstance(valid, AndNode) else [valid]
+    var_dict = {sym_render(var):(var, [var.min, var.max]) for var in valid.vars()}
+
+    for nd in nodes:
+      init_var: Variable = nd.vars()[0]
+      var = var_dict[sym_render(init_var)][1]
+      if isinstance(nd.a, MulNode):
+        if nd.a.b < 0:
+          var[0] = (nd.b // nd.a.b) + 1
+        elif nd.a.b > 0:
+          var[1] = (nd.b // nd.a.b) - 1 if nd.b % nd.a.b == 0 else nd.b // nd.a.b
+      elif isinstance(nd.a, Variable):
+        var[1] = nd.b - 1
+    # We do not allow NumNode because it is constant
+    # TODO: Remove mx != mn
+    sub_dict: dict[Union[Variable, NumNode], Node] = {v:Variable(k, mn, mx) for k, (v, (mn, mx)) in var_dict.items() if mx != mn}
+    valid = valid.substitute(sub_dict)
+    idxy = idxy.substitute(sub_dict)
+
+  b = base_shape[1]
+  idx = (idxy // 4) % b
+  idy = (idxy // (4 * b))
+
+  # Simplify ModNode if possibe
+  if valid.min == 0 and isinstance(idx, ModNode):
+    nds = valid.nodes if isinstance(valid, AndNode) else [valid]
+    ones: List[Node] = []
+    assert isinstance(idx.a, SumNode)
+    idx_nodes = idx.a.flat_components
+    for nd in nds:
+      if not isinstance(idx, ModNode) or not isinstance(nd.a, SumNode): continue
+      nd_flat = nd.a.flat_components
+      if not all(isinstance(node, (MulNode, Variable)) for node in nd_flat): continue
+      nd_vars = nd.vars()
+
+      same, others = partition(idx_nodes, lambda x: (x.a if isinstance(x, MulNode) else x) in nd_vars)
+
+      if not all(isinstance(node, (MulNode, Variable)) for node in same): continue
+
+      if len(same) != len(nd_vars): continue
+
+      first = nd_flat[0]
+      var2 = (first.a, first.b) if isinstance(first, MulNode) else (first, 1)
+
+      for s in same:
+        var1 = (s.a, s.b) if isinstance(s, MulNode) else (s, 1)
+        if var2[0] == var1[0]:
+          k = var1[1] // var2[1]
+          break
+      else:
+        continue
+
+      if Variable.sum(same) != k*(nd.a): continue
+      left_sum = Variable.sum(others)
+      if k < 0:
+        # TODO: More thought on this
+        mnn = max((-nd.b) + 1, min([-lal.b if isinstance(lal, MulNode) else lal for lal in same]))
+        if (left_sum.min + (-k)*mnn) >= b:
+          #ones.append(nd)
+          idx = idx.a - b
+      elif k > 0:
+        if ((nd.b - 1)*k + left_sum.max) < b:
+          #ones.append(nd)
+          idx = idx.a
+      break
+    valid = Variable.ands([i for i in nds if i not in ones])
+
+  # Simplify SumNodes
+  if valid.min == 0:
+    nds = valid.nodes if isinstance(valid, AndNode) else [valid]
+    ones: List[Node] = []
+    flats = []
+    if isinstance(idx, SumNode): flats.append(idx.flat_components)
+    if isinstance(idy, SumNode): flats.append(idy.flat_components)
+    for node in nds:
+      for flat in flats:
+        syms = [i for i in flat if not isinstance(i, NumNode)]
+        sym_sum = Variable.sum(syms)
+        if node.a == sym_sum: ones.append(node)
+        if -(node.a) == sym_sum: ones.append(node)
+
+    valid = Variable.ands([i for i in nds if i not in ones])
+  # This is basically expand with custom variable range
+  if valid.min == 0 and not isinstance(idx, ModNode):
+    variables = tuple(set(valid.vars() + idy.vars() + idx.vars()))
+    val_infer = valid.expand(variables)
+    idx_infer = idx.expand(variables)
+    idy_infer = idy.expand(variables)
+    ones = set()
+    zeros = set()
+    for i, x, y in zip(val_infer, idx_infer, idy_infer):
+      (zeros if i.min == 0 else ones).add((x, y))
+
+    if not ones.intersection(zeros):
+      valid = NumNode(1)
+  """
+  # Simplify sumnodes
+  if valid.min == 0 and not isinstance(idx, ModNode):
+    nds = valid.nodes if isinstance(valid, AndNode) else [valid]
+    ones = []
+    for nd in nds:
+      if not isinstance(nd.a, SumNode): continue
+      for index in (idx, idy):
+        if isinstance(index, SumNode):
+          index_nodes = set(index.flat_components)
+
+          nd_nodes = set(nd.a.flat_components)
+          diff = index_nodes - nd_nodes
+          if len(nd_nodes - index_nodes) == 0:
+            if any(not isinstance(d, (DivNode, NumNode)) for d in diff): break
+            ones.append(nd)
+          else:
+            nd_nodes = set((-nd.a).flat_components)
+            diff = index_nodes - nd_nodes
+            if len(nd_nodes - index_nodes) == 0:
+              if any(not isinstance(d, (DivNode, NumNode)) for d in diff): break
+              ones.append(nd)
+
+    valid = Variable.ands([i for i in nds if i not in ones])
+    
+    def recurse(variables, idx, idy, valid, mem):
+    if len(variables) == 0:
+      if valid in mem: mem[valid].add((idx, idy))
+      else: mem[valid] = set(((idx, idy),))
+    else:
+      var = variables[0]
+      k = 2
+      range_list = list(range(var.min, var.max + 1))
+      range_list = range_list[:k] + range_list[-k:]
+      for i in set(range_list):
+        var_vals = {var:NumNode(i)}
+        val_infer = valid.substitute(var_vals)
+        idx_infer = idx.substitute(var_vals)
+        idy_infer = idy.substitute(var_vals)
+        recurse(variables[1:], idx_infer, idy_infer, val_infer, mem)
+      return mem
+  if False and valid.min == 0 and not isinstance(idx, ModNode):
+    variables = list(set(valid.vars() + idy.vars() + idx.vars()))
+    mem: Dict[NumNode, Any] = {NumNode(1): set(), NumNode(0): set()}
+    mem = recurse(variables, idx, idy, valid, mem)
+    ones = mem[NumNode(1)]
+    zeros = mem[NumNode(0)]
+    if len(ones.intersection(zeros)) == 0:
+      valid = NumNode(1)
+  """
+  if valid.min == 0:
+    print(idx, idy, valid)
+
+  if DEBUG>=5: print("to_image_idx", base_shape, idx.min, idx.max, idy.min, idy.max, idx, idy)
+  return (idx, idy), valid
 
 class UOp(NamedTuple):
   uop: UOps
@@ -129,10 +262,11 @@ def global_load(self, i:int, idxs:Sequence[VariableOrNum], acc=None) -> List[UOp
           buf_uop = self.buf_uops[i]
           assert buf_uop is not None, f"buffer {i} wasn't UOped"
           if isinstance(self.bufs[i].dtype, ImageDType):
-            idx = to_image_idx(self.bufs[i].dtype.shape, idx, valid)
+            idx, valid = to_image_idx(self.bufs[i].dtype.shape, idx, valid)
             rendered_idx = self.uop(UOps.CAST, dtypes._int2, (idx[0].render(self.render_ops, self), idx[1].render(self.render_ops, self)))
           else:
             rendered_idx = idx.render(self.render_ops, self)
+
           if valid.min == 0:
             valid_rendered = valid.render(self.render_ops, self)
             self.load_cache[key] = self.uop(UOps.LOAD, localtype, (buf_uop, rendered_idx, valid_rendered, self.const(invalid_value, localtype)))
@@ -168,7 +302,7 @@ def global_store(self, i:int, idxs:List[VariableOrNum], store:List[UOp]) -> None
     for idx, var in store_offset.items():
       idx, valid = self.sts[i].expr_idxs(idx)
       if isinstance(self.bufs[i].dtype, ImageDType):
-        idx = to_image_idx(self.bufs[i].dtype.shape, idx, valid)
+        idx, valid = to_image_idx(self.bufs[i].dtype.shape, idx, valid)
         rendered_idx = self.uop(UOps.CAST, dtypes._int2, tuple(x.render(self.render_ops, self) for x in idx))
       else:
         rendered_idx = idx.render(self.render_ops, self)
diff --git a/tinygrad/shape/symbolic.py b/tinygrad/shape/symbolic.py
index ae0f8c96b67b5..02d368b13dfd2 100644
--- a/tinygrad/shape/symbolic.py
+++ b/tinygrad/shape/symbolic.py
@@ -97,6 +97,7 @@ def __mod__(self, b:Union[Node,int]):
     assert b > 0
     if b == 1: return NumNode(0)
     if self.min >= 0 and self.max < b: return self
+    if (self.min//b) == (self.max//b): return self - (b*(self.min//b))
     if self.min < 0: return (self - ((self.min//b)*b)) % b
     return create_node(ModNode(self, b))
 


Code review comments:
diff_hunk: @@ -21,25 +21,108 @@ class UOps(Enum):
   LOAD = auto(); STORE = auto(); CONST = auto(); BARRIER = auto() # noqa: E702
   ALU = auto(); WMMA = auto(); CAST = auto(); GEP = auto() # noqa: E702
 
-def to_image_idx(base_shape:Tuple[int, ...], idxy:Node, valid:Node, validhacks=False) -> Tuple[Node, Node]:
-  idy = (idxy//(4*base_shape[1]))
-  if validhacks and valid.min == 0:
-    idx = (idxy//4) + (idy*-base_shape[1])
-    # find the ones in idx that didn't factorize and remove them (TODO: this is not universal)
-    if isinstance(idx, SumNode):
-      unfactored, idx_nodes = partition(idx.nodes, lambda x: isinstance(x, MulNode) and x.b == -base_shape[1])
-      assert len(unfactored) <= 1
-      idx = Variable.sum(idx_nodes)
-      unfactored = (Variable.sum(unfactored) // base_shape[1])
-      idy += unfactored
-      # ugh really...handtuned garbage
-      if idx.min >= (base_shape[1]*3)//4:
-        idx -= base_shape[1]
-        idy += 1
-  else:
-    idx = (idxy//4)%base_shape[1]
-  if DEBUG >= 5: print("to_image_idx", base_shape, idx.min, idx.max, idy.min, idy.max, idx, idy)
-  return idx, idy
+def to_image_idx(base_shape:Tuple[int, ...], idxy:Node, valid:Node) -> Tuple[Tuple[Node, Node], Node]:
+  if valid.min == 0:
+    nodes: List = valid.nodes if isinstance(valid, AndNode) else [valid]
+    var_dict = {sym_render(var):(var, [var.min, var.max]) for var in valid.vars()}
+
+    for nd in nodes:
+      init_var: Variable = nd.vars()[0]
+      var = var_dict[sym_render(init_var)][1]
+      if isinstance(nd.a, MulNode):
+        if nd.a.b < 0:
+          var[0] = (nd.b // nd.a.b) + 1
+        elif nd.a.b > 0:
+          var[1] = (nd.b // nd.a.b) - 1 if nd.b % nd.a.b == 0 else nd.b // nd.a.b
+      elif isinstance(nd.a, Variable):
+        var[1] = nd.b - 1
+    # We do not allow NumNode because it is constant
+    # TODO: Remove mx != mn
+    sub_dict: dict[Union[Variable, NumNode], Node] = {v:Variable(k, mn, mx) for k, (v, (mn, mx)) in var_dict.items() if mx != mn}
+    valid = valid.substitute(sub_dict)
+    idxy = idxy.substitute(sub_dict)
+
+  b = base_shape[1]
+  idx = (idxy // 4) % b
+  idy = (idxy // (4 * b))
path: tinygrad/codegen/linearizer.py
text: @user1:
`idx, idy = (idxy // 4) % base_shape[1], (idxy // (4 * base_shape[1]))`