mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

jit: memplan before compile (#16560)

Browse source
This commit is contained in:
nimlgen 2026-06-10 15:05:15 +03:00 • committed by GitHub
commit 2c9d2c0d31
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
tinygrad/engine/jit.py
View file

@ -69,8 +69,8 @@ def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
# parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index
linear = linear.substitute({u: UOp.param(i, u.dtype, u.shape, u.device) for i,u in enumerate(input_uops)}, walk=True)
linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
linear = memory_plan_rewrite(linear, held_bufs)
linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)
if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View graphed linear")
return linear