import unittest, itertools from tinygrad.codegen.late.coalesce import indexing_simplify from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp, Ops, graph_rewrite from tinygrad.uop.weak import pm_commit_weak from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load from tinygrad.helpers import Context from test.helpers import full_rewrite from test.null.test_uop_symbolic import check_uop_against_string # symbolic-only idx + valid simplification (no late lowering of FLOORDIV/FLOORMOD) def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load, name="simplify_valid_idx") # image-aware idx + valid simplification: adds the codegen-layer matcher that drops provably in-bounds gates def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx") def get_gated_load_uop(valid:UOp, idx:UOp): return UOp.param(0, dtypes.float, 1024).index(idx.valid(valid)).load() def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]): return UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load() def Special(expr, nmax): return UOp.special(nmax, expr) def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax, param=True) def Range(n, nmax): return UOp.range(nmax, n) class TestValidIdxSimplification(unittest.TestCase): def check(self, load, sidx, svalid, extra=()): load = simplify_valid_idx(UOp.sink(load, *extra)).src[0] off = load.src[0].src[1] check_uop_against_string(self, off.get_idx(), sidx) check_uop_against_string(self, off.get_valid(), svalid) def test_cumsum(self): gidx0 = Special("gidx0", 5) lidx0 = Special("lidx0", 4) gate = (gidx0*4+lidx0<19).ne(True) idx = gidx0*4+lidx0-19 load = get_gated_load_uop(gate, idx) self.check(load, "0", "(((lidx0+(gidx0*4))<19)!=True)") def test_simplify_within_valid1(self): ridx0 = Range(0, 4) ridx1 = Range(1, 4) ridx2 = Range(2, 4) ridx3 = Range(3, 4) valid = ((ridx0*3+ridx1)<8) & ((((ridx0*3+ridx1)//8+ridx2*3+ridx3)%4)<2) idx = ridx0+ridx1+ridx2+ridx3 load = get_gated_load_uop(valid, idx) self.check(load, "(((r0+r1)+r2)+r3)", "((((r0*3)+r1)<8)&((((r2*3)+r3)%4)<2))") def test_simplify_within_valid2(self): gidx0 = Special("gidx0", 56) ridx0 = Range(0, 3) alu0 = gidx0+ridx0 valid = (alu0 < 57) & (alu0 >= 1) self.assertIsNone(simplify_valid(valid)) def test_valid_order_matters1(self): ridx0 = Range(0, 2) v0 = ridx0<1 v1 = ((ridx0*5+1)%6)<5 self.assertEqual(simplify_valid(v0&v1).render(), "(r0<1)") self.assertEqual(simplify_valid(v1&v0).render(), "(r0<1)") def test_valid_order_matters2(self): gidx0 = Special("gidx0", 13) gidx1 = Special("gidx1", 13) ridx0 = Range(0, 4) alu0 = (gidx1+(ridx0*13)) v0 = (gidx0+11)%14<11 v1 = (alu0+((gidx0+39)//42))%14<11 v2 = gidx0<3 v3 = alu0<42 for v in itertools.permutations([v0,v1,v2,v3]): self.assertEqual(simplify_valid(v[0]&v[1]&v[2]&v[3]).render(), "False") def test_simplify_valid_from_div(self): x = Variable("x", -100, 100) valid = ((x<0)&((100%x).cast(dtypes.bool))) # NOTE: this simplifies the (100%x) part somehow, still has two clauses self.assertIsNotNone(simplify_valid(valid)) self.assertEqual(len(list(valid.split_uop(Ops.AND))), 2) @unittest.expectedFailure # TODO: fix def test_from_merge_views(self): # taken from test_merges_from_fuzzer1 # generated by # v0 = View(shape=(2, 4), strides=(2, 1), offset=-2, mask=((0, 2), (2, 4)), contiguous=False) # v1 = View(shape=(2, 4, 2, 2), strides=(4, 0, -2, -1), offset=3, mask=None, contiguous=False) # s = ShapeTracker((v0, v1)) # idx, valid = s.to_indexed_uops() # print(f"{idx.render()=}") # print(f"{valid.render()=}") # s = ShapeTracker((View(shape=(2, 4, 2, 2), strides=(2, 0, 0, -1), offset=1, mask=((0, 2), (0, 4), (0, 1), (0, 2)), contiguous=False),)) # idx, valid = s.to_indexed_uops() # print(f"{idx.render()=}") # print(f"{valid.render()=}") ridx0 = Range(0, 2) ridx2 = Range(2, 2) ridx3 = Range(3, 2) idx = (((ridx0*2)+((((ridx2*2)+(ridx3*3))+3)%4))+-2) valid = ((((((ridx2*2)+(ridx3*3))+3)%4)<2)!=True) # noqa: E712 load = get_gated_load_uop(valid, idx) self.check(load, "(((r0*2)+(r3*-1))+1)", "(r2<1)") def test_load_in_valid(self): # from FUSE_ARANGE=1 python test/test_ops.py TestOps.test_scatter_add # can lead to OOB ridx2 = Range(2, 4) lidx0 = Special("lidx0", 3) gidx0 = Special("gidx0", 2) idx=(((lidx0+(gidx0*3))+(ridx2*5))+40) valid = (lidx0+(gidx0*3)) < 5 val7 = get_gated_load_uop(valid, idx) valid2 = valid & val7.cast(dtypes.bool).logical_not() self.assertIsNone(simplify_valid(valid2)) def test_valid_becomes_const1(self): # from DSP mobilenetv2 ridx0 = Range(0, 30) ridx1 = Range(1, 7) ridx2 = Range(2, 2) alu11 = (ridx1+ridx2) alu15 = ((alu11+1)//7) idx = (alu15*-31)+(((((alu11+218)//224)+ridx0)%30)*1568) valid = (ridx2<1)&(ridx1<6) load = get_gated_load_uop(valid, idx) # prevent ridx1 and ridx2 from being shrunk red = load.reduce(ridx1, ridx2, arg=Ops.ADD) self.check(load, "(r0*1568)", "((r2<1)&(r1<6))", extra=(red,)) def test_valid_becomes_const1_z3(self): from z3 import Ints, Solver, And, If, Not, unsat ridx0, ridx1, ridx2, alu11, alu15 = Ints('ridx0 ridx1 ridx2 alu11 alu15') alu11 = (ridx1+ridx2) alu15 = ((alu11+1)/7) idx = (alu15*-31)+(((((alu11+218)/224)+ridx0)%30)*1568) valid = (ridx2<1)&(ridx1<6) load = If(valid, idx, 0) # correct simplification s = Solver() s.add(And(0<=ridx0, ridx0<30, 0<=ridx1, ridx1<7, 0<=ridx2, ridx2<2)) simplifed_idx = (ridx0*1568) simplifed_load = If(valid, simplifed_idx, 0) s.add(Not(load == simplifed_load)) # Check if they are NOT equivalent assert s.check() == unsat, f"The expressions are not equivalent. {s.model()=}" # new solver for a wrong simplified expression s = Solver() s.add(And(0<=ridx0, ridx0<30, 0<=ridx1, ridx1<7, 0<=ridx2, ridx2<2)) wrong_simplifed_idx = (ridx0*1567)+ridx1 wrong_simplifed_load = If(valid, wrong_simplifed_idx, 0) s.add(Not(load == wrong_simplifed_load)) # Check if they are NOT equivalent assert s.check() != unsat, "The expressions are equivalent??" print("The expressions are not equivalent.") print(s.model()) def test_valid_becomes_const2(self): ridx0 = Range(0, 4) ridx1 = Range(1, 4) ridx2 = Range(2, 4) ridx3 = Range(3, 4) # TODO: this should also work without the extra nesting idx = (((ridx0+ridx1)+(ridx2+ridx3)+28)//30) valid = ((ridx0+ridx1)<1).ne(True) & ((ridx2+ridx3)<1).ne(True) load = get_gated_load_uop(valid, idx) self.check(load, "1", "((((r0+r1)<1)!=True)&(((r2+r3)<1)!=True))") def test_valid_with_non_const_rhs(self): ridx0 = Range(0, 1024) ridx1 = Range(1, 4) ridx2 = Range(2, 4) valid = (ridx0<(ridx1*4 + ridx2))&(ridx0<-1).ne(True) idx = ridx0 load = get_gated_load_uop(valid, idx) self.check(load, "r0", "(r0<((r1*4)+r2))") class TestImageSimplification(unittest.TestCase): def check(self, load, svalid, sidx0, sidx1): load = simplify_image_idx(load.sink()).src[0] off = load.src[0] self.assertEqual(len(off.src), 3) idx0, idx1 = off.src[2].get_idx(), off.src[1].get_idx() check_uop_against_string(self, idx0, sidx0) check_uop_against_string(self, idx1, sidx1) self.assertEqual(off.src[1].get_valid(), off.src[2].get_valid()) if svalid is not None: check_uop_against_string(self, off.src[1].get_valid(), svalid) else: self.assertEqual(off.src[1].get_valid(), UOp.const(True), "svalid is None but valid is not True") def test_idx_gt_c(self): # (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid # (idx1 < c+1).ne(True) -> idx > c gidx0 = Special("gidx0", 32) gidx1 = Special("gidx1", 32) shape = (10, 10, 4) load = get_load_image_uop(shape, (gidx1<1).ne(True), (gidx0, gidx1-1)) self.check(load, None, "gidx0", "(gidx1+-1)") load = get_load_image_uop(shape, (gidx1<1).ne(True), (gidx0, gidx1-2)) self.check(load, None, "gidx0", "(gidx1+-2)") # should match any one of the AND clause and drop the matched statement from valid valid = (gidx0<1).ne(True) & (gidx1<1).ne(True) load = get_load_image_uop(shape, valid, (gidx0+1, gidx1-1)) self.check(load, "((gidx0<1)!=True)", "(gidx0+1)", "(gidx1+-1)") valid = (gidx0<1).ne(True) & (gidx1<1).ne(True) load = get_load_image_uop(shape, valid, (gidx0, gidx1-1)) self.check(load, "((gidx0<1)!=True)", "gidx0", "(gidx1+-1)") def test_idx_lt_bound(self): # (idx1 < image_bound) ? (..., idx1) : 0 can drop the valid gidx0 = Special("gidx0", 32) gidx1 = Special("gidx1", 32) load = get_load_image_uop((10, 10, 4), gidx1<10, (gidx0, gidx1)) self.check(load, None, "gidx0", "gidx1") # same thing, valid has a div load = get_load_image_uop((10, 10, 4), gidx1//2<5, (gidx0, gidx1)) self.check(load, None, "gidx0", "gidx1") # 10x20 image, not out of bound load = get_load_image_uop((20, 10, 4), gidx1<10, (gidx0, gidx1)) self.check(load, "(gidx1<10)", "gidx0", "gidx1") def test_generic_idx_lt_bound(self): # (idx1 < image_bound - c) ? (..., idx1 + c) : 0 can drop the valid gidx0 = Special("gidx0", 32) gidx1 = Special("gidx1", 32) shape = (10, 10, 4) load = get_load_image_uop(shape, (gidx1<8), (gidx0, gidx1+2)) self.check(load, None, "gidx0", "(gidx1+2)") load = get_load_image_uop(shape, (gidx1<5), (gidx0, gidx1+5)) self.check(load, None, "gidx0", "(gidx1+5)") def test_valid_empty_set(self): gidx0 = Special("gidx0", 32) gidx1 = Special("gidx1", 32) shape = (32, 32, 4) idx = (gidx0%2, gidx1+2) # not empty load = get_load_image_uop(shape, gidx0<8, idx) self.check(load, "(gidx0<8)", "(gidx0%2)", "(gidx1+2)") # empty -> invalid load = get_load_image_uop(shape, (gidx0<8) & (gidx0<8).ne(True), idx) with Context(NOOPT=1, SPEC=0): load = full_rewrite(load.sink()).src[0] self.assertFalse(load.op_in_backward_slice_with_self(Ops.LOAD)) def test_openpilot_conv1(self): # first conv in openpilot # kernel in tinygrad ae5d1407ee844a97a52ad3756835d38e7e2b9e1b https://gist.github.com/chenyuxyz/39c2d4e9a076b46731c67d345ff066b6 idx1 = Special("idx1", 32) idx2 = Special("idx2", 64) # ridx0 = Variable("ridx0", 0, 5) # ridx1 = Variable("ridx1", 0, 2) # ridx2 = Variable("ridx2", 0, 2) ridx0 = Range(0, 6) ridx1 = Range(1, 3) ridx2 = Range(2, 3) alu1 = ((idx2*2)+ridx1) alu4 = ((idx1*48)+(ridx2*6)+ridx0) valid = ((((idx2*2)+(ridx1))<1).ne(True))&((((idx1*8)+(ridx2))<1).ne(True)) shape = (128, 1536, 4) idx = ((alu4+1530)%1536, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2)) load = get_load_image_uop(shape, valid, idx) self.check(load, None, "((((idx1*48)+(r2*6))+r0)+-6)", "(((idx2*2)+r1)+-1)") def test_openpilot_conv2(self): # conv in test/external/external_test_valid_remove.py idx1 = Special("idx1", 32) idx2 = Special("idx2", 64) # ridx0 = Variable("ridx0", 0, 2) # ridx1 = Variable("ridx1", 0, 2) # ridx2 = Variable("ridx2", 0, 2) ridx0 = Range(0, 3) ridx1 = Range(1, 3) ridx2 = Range(2, 3) alu1 = ((idx2*2)+ridx1) alu3 = ((idx1*24)+(ridx2*3)+ridx0) valid = ((((idx2*2)+ridx1)<1).ne(True))&((((idx1*8)+ridx2)<1).ne(True)) shape = (128, 768, 4) idx = ((alu3+765)%768, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2)) load = get_load_image_uop(shape, valid, idx) self.check(load, None, "((((idx1*24)+(r2*3))+r0)+-3)", "(((idx2*2)+r1)+-1)") def test_openpilot_conv3(self): # in openpilot 0.9.7 idx0 = Special("idx0", 64) idx1 = Special("idx1", 2) idx2 = Special("idx2", 4) ridx0 = Range(0, 7) ridx1 = Range(1, 7) alu2 = ((idx2*2)+ridx0) alu4 = ((idx1*8)+ridx1) alu6 = ((idx1*512)+(ridx1*64)+idx0) valid = (alu2<11)&(alu4<3).ne(True) shape = (8, 1024, 4) idx = (((alu6+832)%1024),(alu2+((idx1+((ridx1+5)//8)+1)//2)+(-4))) load = get_load_image_uop(shape, valid, idx) self.check(load, "(((idx2*2)+r0)<11)", "(idx0+(idx1*512+r1*64)+-192)", "((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)") def test_simplify1(self): # idx has the form (A % m, A // m + k) and valid has (c0 < A) and (A < c1) gidx = Special("gidx", 512) valid = (gidx<488) & (gidx<480).ne(True) idx = ((gidx*3+18)%26, (gidx*3+18)//26-56) load = get_load_image_uop((1, 26, 4), valid, idx) self.check(load, None, "((gidx*3)+-1438)", "0") def test_simplify2(self): # from DEV=CL DEBUG=4 FORWARD_ONLY=1 IMAGE=2 python3 test/test_ops.py TestOps.test_simple_padding_conv2d lidx = Special("lidx", 4) valid = (lidx<3) & (lidx<1).ne(True) idx = ((lidx+1)%2, (lidx+1)//2-1) load = get_load_image_uop((1, 2, 4), valid, idx) self.check(load, None, "(lidx+-1)", "0") def test_simplify3(self): # from openpilot idx0 = Special("idx0", 265) valid = (idx0<201).ne(True) idx = ((idx0+55)%64, (idx0+55)//64-4) load = get_load_image_uop((1, 64, 4), valid, idx) self.check(load, None, "(idx0+-201)", "0") def test_simplify4(self): idx0 = Special("idx0", 512) shape = (4, 64, 4) alu2 = ((idx0*4+1)%32) alu3 = ((idx0*4+2)%32) alu4 = ((idx0*4+3)%32) alu5 = (idx0*4%32) alu8 = (idx0//8%32//4) alu9 = idx0<256 load = get_load_image_uop(shape, alu9, (((alu8+(alu2*8))%64),(alu2//8))) self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+8)", "(idx0//2%4)") load = get_load_image_uop(shape, alu9, (((alu8+(alu3*8))%64),(alu3//8))) self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+16)", "(idx0//2%4)") load = get_load_image_uop(shape, alu9, (((alu8+(alu4*8))%64),(alu4//8))) self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+24)", "(idx0//2%4)") load = get_load_image_uop(shape, alu9, (((alu8+(alu5*8))%64),(alu5//8))) self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32)", "(idx0//2%4)") def test_simplify5(self): # openpilot 0.9.7, chunk replacement to simplify shape = (10, 384, 4) idx0 = Special("idx0", 16) idx1 = Special("idx1", 24) alu0 = idx0*4 alu1 = (idx1*256)+alu0 alu2 = idx1//3 alu3 = ((alu1+1)%768) idx = ((idx0+((((alu3//640)+alu2)%8)*16)+128),((alu3//64)%10)) valid = alu3<640 load = get_load_image_uop(shape, valid, idx) self.check(load, None, "((idx0+((idx1//3)*16))+128)", "((idx1%3)*4)") def test_simplify6(self): # from openpilot # the valid implies the numerator of the div/mod is positive and can be simplified with floordiv rules idx1 = Special("idx1", 16) idx2 = Special("idx2", 64) ridx3 = Range(3, 3) ridx4 = Range(4, 3) ridx5 = Range(5, 3) alu0 = ((idx2*1536)+(ridx4*768)+ridx3+(idx1*24)+(ridx5*3)+-771)%768 alu1 = ((idx2*1536)+(ridx4*768)+ridx3+(idx1*24)+(ridx5*3)+-771)//768 valid = (((idx2+ridx4)<1)!=1)&(((idx1+ridx5)<1)!=1) load = get_load_image_uop((128, 768, 4), valid, (alu0, alu1)) self.check(load, None, "((((idx1*24)+(r5*3))+r3)+-3)", "(((idx2*2)+r4)+-1)") def test_simplify7(self): # DEBUG=2 ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 # noqa: E501 # kernel 143 gidx0 = Special("gidx0", 32) lidx0 = Special("lidx0", 16) lidx1 = Special("lidx1", 8) r0 = Range(0, 7) # buf.render()='UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), arg=1, src=())' alu0 = ((gidx0*2+(lidx0*128+r0*64+lidx1*8+-183)%64*64+(lidx0*128+r0*64+lidx1*8+-183)//64%32*4096+1)//4%1024) alu1 = ((gidx0*2+(lidx0*128+r0*64+lidx1*8+-183)%64*64+(lidx0*128+r0*64+lidx1*8+-183)//64%32*4096+1)//4096) valid = ((lidx1<7)&((((lidx0*2+r0)<3)!=1)&((lidx0*2+r0)<35))) load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1)) self.check(load, None, "(lidx1*128+gidx0//2+144)", "(lidx0*2+r0+-3)") # same idx, written without the inline simplification of the inner div/mod alu0 = ((gidx0*2+lidx1*512+(lidx0*8192+r0*4096)+-11711)//4%1024) alu1 = (lidx0*2+r0+-3) valid = ((lidx1<7)&((((lidx0*2+r0)<3)!=1)&((lidx0*2+r0)<35))) load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1)) self.check(load, None, "(lidx1*128+gidx0//2+144)", "(lidx0*2+r0+-3)") def test_simplify8(self): # from openpilot compile3, kernel r_4_16_8_16_4_4_3_3n1 # valid guarantees A >= 0, so divmod simplifies and gate is removed gidx0 = Special("gidx0", 16) gidx1 = Special("gidx1", 4) lidx0 = Special("lidx0", 8) lidx1 = Special("lidx1", 16) A = gidx0 + gidx1*8192 + lidx0*1024 + lidx1*64 - 1040 valid = ((lidx1 < 1).ne(True)) & (((gidx1 + lidx0) < 1).ne(True)) load = get_load_image_uop((32, 1024, 4), valid, (A % 1024, A // 1024)) self.check(load, None, "(gidx0+lidx1*64+-16)", "(lidx0+gidx1*8+-1)") def test_simplify9(self): # from openpilot compile3, kernel r_32_16_8_4_4_7_7 (image 1x16384) # valid guarantees A1 >= 0 and A1 < 512, gate should be removable gidx0 = Special("gidx0", 32) lidx0 = Special("lidx0", 16) lidx1 = Special("lidx1", 8) r0 = Range(0, 7) A1 = lidx0*32 + r0*32 + lidx1*4 - 99 valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 19) alu0 = gidx0 + (A1 % 32)*32 + (A1 // 32 % 16)*1024 load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(0))) try: self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0") except AssertionError: # TODO: fold valid self.check(load, "(((lidx1<1)!=True)&((lidx0+r0)<19))", "(gidx0+lidx1*128+(lidx0*1024+r0*1024)+-3168)", "0") def test_simplify10(self): # from openpilot compile3, kernel r_16_8_4_4_4_4_7_7 (image 1x8192) # valid guarantees A1 >= 0 and A1 < 128, gate should be removable gidx0 = Special("gidx0", 16) lidx0 = Special("lidx0", 8) lidx1 = Special("lidx1", 4) lidx2 = Special("lidx2", 4) r0 = Range(0, 7) A1 = lidx0*16 + r0*16 + lidx1*4 - 51 valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 11) alu0 = lidx2 + gidx0*4 + (A1 % 16)*64 + (A1 // 16 % 8)*1024 load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(0))) try: self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0") except AssertionError: # TODO: fold valid self.check(load, "(((lidx1<1)!=True)&((lidx0+r0)<11))", "(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0") def test_drop_non_monotonic_window(self): # two-sided window valid (645 <= gidx0 < 653) on a non-monotonic index (lane split via %4 and //4): # gidx0 outside the window pushes idx_x out of the (1, 48) image, so the gate is dropped gidx0 = Special("gidx0", 1064) r12 = Range(12, 3) valid = ((gidx0 < 645).ne(True)) & (gidx0 < 653) idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(0)) load = get_load_image_uop((1, 48, 4), valid, idx) self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0") def test_drop_gate_committed_in_the_index_pass(self): # the fused index pass runs without symbolic, so committing a weak src must not leave a CAST that # symbolic later folds inside the index only: the gate's copy of the expression has to stay the same node f = UOp.variable("f", 0.0, 9.0, dtypes.float) idx_y = (f + UOp.const(1.0)).cast(dtypes.int) load = get_load_image_uop((10, 10, 4), (UOp.const(-1) < idx_y) & (idx_y < UOp.const(10)), (Special("gidx0", 10), idx_y)) off = graph_rewrite(load.sink(), pm_commit_weak+indexing_simplify).src[0].src[0] self.assertEqual(off.src[1].get_valid(), UOp.const(True)) class TestDropTrueGate(unittest.TestCase): def test_drop_true_gate_on_index(self): # test that INDEX with a constant True valid gets simplified to drop the valid from tinygrad.codegen.late.coalesce import indexing_simplify from tinygrad.uop.ops import graph_rewrite from tinygrad.uop.symbolic import sym buf = UOp.param(0, dtypes.int, 1) idx = UOp.const(0) true_gate = UOp.const(True) index_with_gate = buf.index(idx.valid(true_gate)) # apply the optimization result = graph_rewrite(index_with_gate, sym+indexing_simplify) # the True valid should be dropped (INDEX should only have 2 sources) self.assertEqual(len(result.src), 2, "True valid should be dropped from INDEX") class TestRangeShrink(unittest.TestCase): def get_ranges(self, sink): with Context(NOOPT=1, SPEC=0): result = full_rewrite(sink) return [u for u in result.toposort() if u.op is Ops.RANGE] def assert_range_end(self, ranges:list[UOp], end:int): self.assertEqual(len(ranges), 1) with Context(NOOPT=1, SPEC=0): expected = full_rewrite(UOp.const(end, dtypes.int).sink()).src[0] self.assertIs(ranges[0].src[0], expected) def test_range_shrink_single_guard(self): # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3 r = Range(0, 204) load = get_gated_load_uop(r < UOp.const(4), r) ranges = self.get_ranges(load.sink()) self.assert_range_end(ranges, 4) def test_range_shrink_picks_max_guard(self): # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8 r = Range(0, 204) load1 = get_gated_load_uop(r < UOp.const(4), r) load2 = get_gated_load_uop(r < UOp.const(8), r) ranges = self.get_ranges(UOp.sink(load1, load2)) self.assert_range_end(ranges, 8) def test_range_no_shrink_guard_ge_max(self): # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain) r = Range(0, 204) load = get_gated_load_uop(r < UOp.const(300), r) ranges = self.get_ranges(load.sink()) self.assert_range_end(ranges, 204) def test_range_no_shrink_when_unguarded_elsewhere(self): # one load guards r < 4, but another load uses r without a gate -> no shrink r = Range(0, 204) load1 = get_gated_load_uop(r < UOp.const(4), r) load2 = UOp.param(1, dtypes.float, 204).index(r).load() ranges = self.get_ranges(UOp.sink(load1, load2)) self.assert_range_end(ranges, 204) def test_range_no_shrink_when_used_in_reduce(self): # range used in both a gated load AND directly in the reduce expression -> no shrink r = Range(0, 204) gated_load = get_gated_load_uop(r < UOp.const(4), r) red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD) ranges = self.get_ranges(red.sink()) self.assert_range_end(ranges, 204) def test_range_shrink_to_single_iteration(self): # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely r = Range(0, 204) load = get_gated_load_uop(r < UOp.const(1), r) ranges = self.get_ranges(load.sink()) self.assertEqual(len(ranges), 0) def test_range_shrink_store_where_invalid(self): # emulates mask.where(x.pad_to(mask.shape), Invalid): range should shrink accordingly from tinygrad.dtype import Invalid r = Range(0, 204) x = (r < 4).where(UOp.const(1.0), Invalid) ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r < 4).where(x, Invalid)).sink()) self.assert_range_end(ranges, 4) def test_range_shrink_store_where_invalid_flipped(self): # above, but flipped from tinygrad.dtype import Invalid r = Range(0, 204) x = (r < 4).where(UOp.const(1.0), Invalid) ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r >= 4).where(Invalid, x)).sink()) self.assert_range_end(ranges, 4) if __name__ == '__main__': unittest.main()