diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 084c110d99..a010d082fb 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -44,6 +44,10 @@ try: UOp(Ops.NOOP, arg=z3.Bool(f"float_cmp{ctx[1].setdefault(x, len(ctx[1]))}",ctx=ctx[0].ctx))), ]) + def uops_to_z3(solver, *uops: UOp) -> 'list[z3.ExprRef]': + with Context(TRACK_MATCH_STATS=0): # cant pickle z3 objects + return [s.arg for s in graph_rewrite(uops[0].sink(*uops[1:]), z3_renderer, ctx=(solver, {})).src] + z3_imported = True except (ImportError, AttributeError): z3_imported = False @@ -124,9 +128,8 @@ def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)): if not z3_imported: raise ImportError("z3 is required for bounds checking, try IGNORE_OOB=0 or \"pip install z3-solver\"") solver = z3.Solver(ctx=z3.Context()) - z3_sink = graph_rewrite(idx.src[1].sink(mask), z3_renderer, ctx=(solver, {})) - z3_idx = z3_sink.src[0].arg - solver.add(z3_sink.src[1].arg) + z3_idx, z3_mask = uops_to_z3(solver, idx.src[1], mask) + solver.add(z3_mask) if solver.check((z3_idx<0)|(sz<=z3_idx)) == z3.sat: print(f"idx={idx.src[1].render(simplify=False)}") print(f"mask & gate={mask.render(simplify=False)}") diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 196f89c8bb..98fea4bc71 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -80,7 +80,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]: if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None: label += f"\n{shape_to_str(u.shape)}" elif len(rngs:=u.ranges): - label += f"\n({','.join(sorted([colored(str(x.arg[0]), axis_colors[x.arg[1]]) for x in rngs]))})" + label += f"\n({','.join([colored(str(x.arg[0]), axis_colors[x.arg[1]]) for x in sorted(rngs, key=lambda x: x.arg[0])])})" except Exception: label += "\n" if (ref:=ref_map.get(u.arg.ast) if u.op is Ops.KERNEL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}"