nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

fix nolocals and beam (#16232)

Browse Source
This commit is contained in:
nimlgen
2026-05-16 18:09:19 +03:00
committed by GitHub
parent ebcb7b7cc0
commit bef5f717bc
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/engine/realize.py
+3 -2
View File
@@ -243,9 +243,9 @@ if getenv("HCQ2"):
from extra.hcq2.hcq2 import pm_hcq_exec
pm_exec = pm_hcq_exec + pm_exec
def compile_linear(linear:UOp, beam=0, validate=False) -> UOp:
def compile_linear(linear:UOp, beam:int|None=None, validate=False) -> UOp:
if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
if (beam_val:=(beam or BEAM.value)) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
@@ -260,4 +260,5 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None
else:
from tinygrad.tensor import Tensor
with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
call = compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0).src[0]
return cast(float, pm_exec.rewrite(call, ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)))