diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 58f95a5d34..80bffd66f5 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -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)