mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-06-24 02:14:17 +00:00
4d6e407eb0* fast idiv for signed ints * Add rule and test * fix tests * redo fuzz_fast_idiv to do negative ints as well * adjust comments * remove unused imports
26 lines
960 B
Python
26 lines
960 B
Python
import random
|
|
import z3
|
|
from tinygrad import dtypes
|
|
from tinygrad.uop.spec import z3_renderer, z3_cdiv
|
|
from tinygrad.uop.ops import UOp, graph_rewrite
|
|
from tinygrad.uop.transcendental import fast_idiv
|
|
random.seed(42)
|
|
|
|
powers_of_two = [2**i for i in range(64)]
|
|
if __name__ == "__main__":
|
|
for i in range(10_000):
|
|
if i % 1000 == 0:
|
|
print(f"Progress: {i}")
|
|
dt = random.choice(dtypes.ints)
|
|
u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
|
|
d = random.randint(1, max(1, u.arg[2]))
|
|
if d in powers_of_two: continue
|
|
expr = fast_idiv(None, u, d)
|
|
if expr is None: continue
|
|
|
|
solver = z3.Solver()
|
|
z3_sink = graph_rewrite(expr.sink(u), z3_renderer, ctx=(solver, {}))
|
|
z3_expr, x = z3_sink.src[0].arg, z3_sink.src[1].arg
|
|
|
|
if solver.check(z3_expr != z3_cdiv(x, d)) == z3.sat:
|
|
assert False, f"Failed: {expr.render()} != x//{d} at x={solver.model()}\nx={u}\nd={d}\n{z3_expr=}\n{x/d=}"
|