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

minor cleanup in jit (#4989)

Browse Source 
found a non-deterministic bug in jit with multiple variables. but first cleanup some variable names.
[run_process_replay]
This commit is contained in:
chenyu
2024-06-15 23:43:17 -04:00
committed by GitHub
parent 44dfa37c70
commit 5eb8001514
1 changed files with 29 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/engine/jit.py
+29 -27
View File
@@ -134,16 +134,17 @@ class TinyJit(Generic[ReturnType]):
def __call__(self, *args, **kwargs) -> ReturnType:
input_tensors: List[Tuple[Union[int, str], Tensor]] = \
[(cast(Union[int, str], k),v) for k,v in itertools.chain(enumerate(args), sorted(kwargs.items())) if v.__class__ is Tensor]
if len(input_tensors): Tensor.realize(*[x[1] for x in input_tensors])
lbs: List[LazyBuffer] = flatten([v.lazydata.lbs for _,v in input_tensors])
expected_sts_var_dtype_device = [(*x.st.unbind(), x.dtype, x.device) for x in lbs]
input_rawbuffers: List[Buffer] = [v.base.realized for v in lbs if v.base.realized is not None]
assert len(set(input_rawbuffers)) == len(input_rawbuffers), "duplicate inputs to JIT"
var_vals: Dict[Variable, int] = merge_dicts([x[1] for x in expected_sts_var_dtype_device] + \
[dict(x.unbind() for x in itertools.chain(args, kwargs.values()) if isinstance(x, Variable))])
expected_names, expected_lbs = [x[0] for x in input_tensors], [(x[0], tuple(x[1].keys()), x[2], x[3]) for x in expected_sts_var_dtype_device]
[(cast(Union[int, str], name),t) for name,t in itertools.chain(enumerate(args), sorted(kwargs.items())) if t.__class__ is Tensor]
if input_tensors: Tensor.realize(*[t for _,t in input_tensors])
names: List[Union[int, str]] = [name for name,_ in input_tensors]
lbs: List[LazyBuffer] = flatten([t.lazydata.lbs for _,t in input_tensors])
st_varvals_dtype_device = [(*lb.st.unbind(), lb.dtype, lb.device) for lb in lbs]
input_buffers: List[Buffer] = [lb.base.realized for lb in lbs if lb.base.realized is not None]
assert len(set(input_buffers)) == len(input_buffers), "duplicate inputs to JIT"
var_vals: Dict[Variable, int] = merge_dicts([varvals for _,varvals,_,_ in st_varvals_dtype_device] + \
[dict(v.unbind() for v in itertools.chain(args, kwargs.values()) if isinstance(v, Variable))])
# TODO: var here is not sorted
st_vars_dtype_device = [(x[0], tuple(x[1].keys()), x[2], x[3]) for x in st_varvals_dtype_device]
if self.cnt == 0:
# jit ignore
with Context(BEAM=0 if getenv("IGNORE_JIT_FIRST_BEAM") else BEAM.value):
@@ -151,8 +152,8 @@ class TinyJit(Generic[ReturnType]):
if len(params:=get_parameters(self.ret)): Tensor.realize(params[0], *params[1:])
elif self.cnt == 1:
# jit capture
self.expected_names: List[Union[int, str]] = expected_names
self.expected_lbs: List[Tuple[ShapeTracker, Tuple[Variable, ...], DType, str]] = expected_lbs
self.expected_names: List[Union[int, str]] = names
self.expected_st_vars_dtype_device: List[Tuple[ShapeTracker, Tuple[Variable, ...], DType, str]] = st_vars_dtype_device
with Context(GRAPH=getenv("JITGRAPH", GRAPH.value), BEAM=getenv("JITBEAM", BEAM.value)):
capturing.append(self)
self.ret = self.fxn(*args, **kwargs)
@@ -160,31 +161,32 @@ class TinyJit(Generic[ReturnType]):
capturing.clear()
del self.buffer_replace
assert len(self.jit_cache), "didn't JIT anything!"
if DEBUG >= 1: print(f"JIT captured {len(self.jit_cache)} kernels with {len(input_rawbuffers)} inputs")
if DEBUG >= 1: print(f"JIT captured {len(self.jit_cache)} kernels with {len(input_buffers)} inputs")
# track inputs that are views of buffers
for ji in self.jit_cache:
for b in ji.bufs:
if b is not None and b._base is not None and b._base in input_rawbuffers:
input_rawbuffers.append(b)
self.extra_view_inputs.append((input_rawbuffers.index(b.base), b.offset, b.device, b.size, b.dtype))
for item in self.jit_cache:
for b in item.bufs:
if b is not None and b._base is not None and b._base in input_buffers:
input_buffers.append(b)
self.extra_view_inputs.append((input_buffers.index(b.base), b.offset, b.device, b.size, b.dtype))
# memory planning (optional)
assigned = _internal_memory_planner([cast(List[Buffer], x.bufs) for x in self.jit_cache], debug_prefix="JIT ")
self.jit_cache = [ExecItem(ei.prg, [assigned.get(x,x).ensure_allocated() for x in ei.bufs if x is not None]) for ei in self.jit_cache]
assigned = _internal_memory_planner([cast(List[Buffer], item.bufs) for item in self.jit_cache], debug_prefix="JIT ")
self.jit_cache = [ExecItem(item.prg, [assigned.get(b,b).ensure_allocated() for b in item.bufs if b is not None]) for item in self.jit_cache]
# Condense the items into a graph executor.
if JIT < 2: self.jit_cache = apply_graph_to_jit(self.jit_cache, input_rawbuffers, var_vals)
if JIT < 2: self.jit_cache = apply_graph_to_jit(self.jit_cache, input_buffers, var_vals)
self.input_replace = get_input_replace(self.jit_cache, input_rawbuffers)
if DEBUG >= 1 and len(set(self.input_replace.values())) != len(input_rawbuffers): print("WARNING: some input tensors not found")
self.input_replace = get_input_replace(self.jit_cache, input_buffers)
if DEBUG >= 1 and len(set(self.input_replace.values())) != len(input_buffers): print("WARNING: some input tensors not found")
elif self.cnt >= 2:
# jit exec
assert self.expected_names == expected_names, f"args mismatch in JIT: {self.expected_names=} != {expected_names}"
assert self.expected_lbs == expected_lbs, f"args mismatch in JIT: {self.expected_lbs=} != {expected_lbs=}"
assert self.expected_names == names, f"args mismatch in JIT: {self.expected_names=} != {names}"
assert self.expected_st_vars_dtype_device == st_vars_dtype_device, \
f"args mismatch in JIT: {self.expected_st_vars_dtype_device=} != {st_vars_dtype_device=}"
for idx, offset, device, size, dtype in self.extra_view_inputs:
input_rawbuffers.append(Buffer(device, size, dtype, base=input_rawbuffers[idx], offset=offset).ensure_allocated())
for (j,i),input_idx in self.input_replace.items(): self.jit_cache[j].bufs[i] = input_rawbuffers[input_idx]
input_buffers.append(Buffer(device, size, dtype, base=input_buffers[idx], offset=offset).ensure_allocated())
for (j,i),input_idx in self.input_replace.items(): self.jit_cache[j].bufs[i] = input_buffers[input_idx]
if DEBUG >= 1 and len(self.jit_cache) >= 10: print(f"jit execs {len(self.jit_cache)} kernels")
for ei in self.jit_cache: ei.run(var_vals, jit=True)