sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions Activity

increase jit batch size progressivly (#3771)

Browse Source 
Co-authored-by: chenyu <[email protected]>
This commit is contained in:
nimlgen
2024-03-16 11:58:11 -04:00
committed by GitHub
co-authored by chenyu
parent 77febb44e6
commit 987a055c0d
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/features/jit.py
+4 -2
View File
@@ -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()