nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

undo loaded comparison swap (#12436)

Browse Source 
* add rule

* add a test
This commit is contained in:
Sieds Lykles
2025-10-03 06:57:29 +02:00
committed by GitHub
parent f203d8b221
commit 0047bcc535
2 changed files with 11 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_uop_graph.py
+8
View File
@@ -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)
tinygrad/codegen/late/devectorizer.py
+3
View File
@@ -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 *****