From 987a055c0dbb5ff077b5e79ea2b8abfd0cf4e1be Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Sat, 16 Mar 2024 18:58:11 +0300 Subject: [PATCH] increase jit batch size progressivly (#3771) Co-authored-by: chenyu --- tinygrad/features/jit.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tinygrad/features/jit.py b/tinygrad/features/jit.py index b6c9dc12cd..35834fb362 100644 --- a/tinygrad/features/jit.py +++ b/tinygrad/features/jit.py @@ -36,15 +36,17 @@ def get_jc_idxs_with_updatable_var_vals(jit_cache: List[JitItem]) -> List[int]: def apply_graph_to_jit(jit_cache: List[JitItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]) -> List[JitItem]: # Split JIT cache into batches for faster graph execution. # This allows the accelerator to run some batches while subsequent graphs are still being updated. + max_batch_size = getenv("JIT_BATCH_SIZE", 32) graphed_jit_cache: List[JitItem] = [] current_batch: List[JitItem] = [] current_device: Optional[Compiled] = None def flush_batch(): - nonlocal current_batch, current_device + nonlocal current_batch, current_device, max_batch_size try: if len(current_batch) <= 1 or current_device is None: raise GraphException("only one kernel doesn't graph") graphed_jit_cache.append(JitItem(current_device.graph(current_batch, input_rawbuffers, var_vals), cast(List[Optional[Buffer]], input_rawbuffers))) # noqa: E501 + max_batch_size *= 2 if DEBUG >= 2: print(f"\tJIT GRAPHing batch with {len(current_batch)} kernels on device {current_device}") except GraphException as e: graphed_jit_cache.extend(current_batch) @@ -58,7 +60,7 @@ def apply_graph_to_jit(jit_cache: List[JitItem], input_rawbuffers: List[Buffer], elif isinstance(ji.prg, BufferXfer) and ji.rawbufs[0] and ji.rawbufs[0].d.dname.startswith("HSA"): ji_graph_dev = ji.rawbufs[0].d can_be_graphed = ji_graph_dev and ji_graph_dev.graph - can_extend_graph_batch = can_be_graphed and len(current_batch) < getenv("JIT_BATCH_SIZE", 64) and (ji_graph_dev == current_device or + can_extend_graph_batch = can_be_graphed and len(current_batch) < max_batch_size and (ji_graph_dev == current_device or (isinstance(ji_graph_dev.graph, type) and issubclass(ji_graph_dev.graph, MultiDeviceJITGraph) and type(ji_graph_dev) == type(current_device))) #type:ignore if not can_extend_graph_batch and len(current_batch) > 0: flush_batch()