diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py index 9f6d5c5ccb..5ea2302d32 100644 --- a/test/test_uop_graph.py +++ b/test/test_uop_graph.py @@ -461,6 +461,8 @@ class TestUOpGraph(unittest.TestCase): if u.op is Ops.STORE: assert u.src[1].arg==5 def test_load_idx_becomes_int(self): + # These loads wont overflow int since we know from the gate that the value is bounded + r0 = UOp.range(10, 0) d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 1) l0 = UOp(Ops.LOAD, dtypes.long, (d0.index(UOp.const(dtypes.int, 0)),)).cast(dtypes.index) @@ -471,6 +473,12 @@ class TestUOpGraph(unittest.TestCase): for u in uops: if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int) + valid = (10*r0<5-l0).ne(True)&(l0<3000) + l2 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx.valid(valid)),)) + uops = to_uops_list([l2]) + for u in uops: + if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int) + def test_in_out_of_bounds_access(self): with Context(IGNORE_OOB=0): glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0) diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 433a46c64c..50a33b5ffe 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -50,6 +50,7 @@ def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp, # remove the gate from the index return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val, *store.src[2:]) +def no_load(u:UOp) -> bool: return not any(x.op is Ops.LOAD for x in u.sparents) load_store_indexing = PatternMatcher([ # image load valid idx simplification (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)), @@ -60,6 +61,8 @@ load_store_indexing = PatternMatcher([ # delete_redundant_gates (after expand) (UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")), UPat.var("val")), name="store", allow_any_len=True), delete_redundant_gates), + # we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes a pattern in reduce_collapse + (UPat.var("c")<(UPat.var("x", dtypes.index)+UPat.var("y")), lambda x,y,c: (-x < -(c-y)) if no_load(y) and no_load(c) and not no_load(x) else None), ]) # ***** load/store grouping *****