You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
33 lines
1.6 KiB
33 lines
1.6 KiB
import random
|
|
from z3 import Int, Solver, sat
|
|
from tinygrad import dtypes, Device
|
|
from tinygrad.ops import UOp, Ops, UPat, graph_rewrite, PatternMatcher
|
|
from tinygrad.codegen.devectorizer import fast_idiv
|
|
random.seed(42)
|
|
|
|
z3_renderer = PatternMatcher([
|
|
(UPat((Ops.DEFINE_VAR, Ops.SPECIAL), name="x"), lambda x: UOp(Ops.NOOP, arg=x.arg[0])),
|
|
# Because fast_idiv only works for non-negative integers we can emulate machine arithmetic with modulo operations.
|
|
(UPat(Ops.SHR, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(({x.src[0].arg}/(2**{x.src[1].arg}))%{dtypes.max(x.dtype)+1})")),
|
|
(UPat(Ops.MUL, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(({x.src[0].arg}*{x.src[1].arg})%{dtypes.max(x.dtype)+1})")),
|
|
(UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: UOp(Ops.NOOP, arg=str(x.arg))),
|
|
(UPat(Ops.CAST, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}")),
|
|
])
|
|
|
|
def render(self) -> str:
|
|
ret = graph_rewrite(self.simplify(), z3_renderer)
|
|
return ret.arg if ret.op is Ops.NOOP else str(ret)
|
|
|
|
if __name__ == "__main__":
|
|
x = Int('x')
|
|
for _ in range(10_000):
|
|
dt = random.choice(dtypes.ints)
|
|
u = UOp(Ops.DEFINE_VAR, dt, arg=('x', 0, random.randint(1, dtypes.max(dt))), src=())
|
|
d = random.randint(1, max(1, u.arg[2]))
|
|
|
|
expr = fast_idiv(Device[Device.DEFAULT].renderer, u, d)
|
|
if expr is None: continue
|
|
solver = Solver()
|
|
solver.add(x>=u.arg[1], x<=u.arg[2])
|
|
if solver.check(eval(render(expr)) != x/d) == sat:
|
|
assert False, f"Failed: {render(expr)} != x//{d} at x={solver.model()[x]}\nx={u}\nd={d}"
|
|
|