mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-12 15:33:45 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
595 lines
24 KiB
Python
595 lines
24 KiB
Python
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_bitwise_and_is_not_a_valid(self):
|
|
ridx0 = Range(0, 16)
|
|
self.assertEqual(simplify_valid_idx(UOp.sink((ridx0 & UOp.const(12, dtypes.int)) & ridx0)).src[0].render(), "((int)(r0)&12&(int)(r0))")
|
|
|
|
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()
|