github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix wait in jit (#2408)

Browse Source
This commit is contained in:
nimlgen
2023-11-24 00:54:13 +03:00
committed by GitHub
parent b927942d58
commit 397c093656
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tinygrad/jit.py
View File

@@ -62,7 +62,7 @@ class TinyJit(Generic[ReturnType]):
assert self.expected_vals == expected_vals, "mismatch of var_vals"
assert self.expected_name_sts_dtype == expected_name_sts_dtype, f"mismatch of sts, expected {self.expected_name_sts_dtype} got {expected_name_sts_dtype}"
for (j,i),input_idx in self.input_replace.items(): self.jit_cache[j].rawbufs[i] = input_rawbuffers[input_idx]
for ji in self.jit_cache: ji.prg(cast(List[RawBuffer], ji.rawbufs), var_vals, jit=True)
for ji in self.jit_cache: ji.prg(cast(List[RawBuffer], ji.rawbufs), var_vals, wait=DEBUG>=2, jit=True)
elif self.cnt == 1:
# jit capture
self.expected_vals, self.expected_name_sts_dtype = expected_vals, expected_name_sts_dtype