forked from tinygrad/tinygrad
optimize symbolic-related updates in graphs (#5727)
* try * faster * cleaner * better? * better? * cleaner * fixes * unused * mypy * fix clang * remove comment * better var names * rename * fix cuda * rename
This commit is contained in:
nimlgen
+26
-9
@@ -3,11 +3,11 @@ from typing import TypeVar, Generic, Callable, List, Tuple, Union, Dict, cast, O
|
||||
import functools, itertools, collections
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.lazy import LazyBuffer
|
||||
from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, GRAPH, BEAM, getenv, all_int, GraphException, colored, JIT
|
||||
from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, GRAPH, BEAM, getenv, all_int, GraphException, colored, JIT, dedup
|
||||
from tinygrad.device import Buffer, Compiled, Device
|
||||
from tinygrad.dtype import DType
|
||||
from tinygrad.shape.shapetracker import ShapeTracker
|
||||
from tinygrad.shape.symbolic import Variable, sint
|
||||
from tinygrad.shape.symbolic import Variable, sint, sym_infer
|
||||
from tinygrad.engine.realize import ExecItem, capturing, EmptyOp, ViewOp, BufferXfer, CompiledRunner, Runner
|
||||
from tinygrad.engine.schedule import _internal_memory_planner
|
||||
from tinygrad.nn.state import get_parameters
|
||||
@@ -70,24 +70,41 @@ def get_input_replace(jit_cache: List[ExecItem], input_rawbuffers:List[Buffer])
|
||||
class GraphRunner(Runner): # pylint: disable=abstract-method
|
||||
def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
|
||||
self.jit_cache = jit_cache
|
||||
self.input_replace = get_input_replace(jit_cache, input_rawbuffers)
|
||||
self.jc_idx_with_updatable_launch_dims = []
|
||||
self.jc_idx_with_updatable_var_vals = []
|
||||
self.input_replace:Dict[Tuple[int, int], int] = get_input_replace(jit_cache, input_rawbuffers)
|
||||
self.var_vals_replace:Dict[int, List[int]] = {}
|
||||
self.launch_dims_replace:Dict[int, Tuple[Optional[int], Optional[int]]] = {}
|
||||
|
||||
op_estimate: sint = 0
|
||||
mem_estimate: sint = 0
|
||||
lds_estimate: sint = 0
|
||||
|
||||
self.vars = sorted(var_vals.keys(), key=lambda v: v.expr)
|
||||
self.symbolic_dims = dedup([tuple(d) for ji in jit_cache if isinstance(ji.prg, CompiledRunner) and (d:=ji.prg.p.local_size) and not all_int(d)] +
|
||||
[tuple(d) for ji in jit_cache if isinstance(ji.prg, CompiledRunner) and (d:=ji.prg.p.global_size) and not all_int(d)])
|
||||
def find_symbolic_dim(dim): return self.symbolic_dims.index(tuple(dim)) if dim is not None and tuple(dim) in self.symbolic_dims else None
|
||||
|
||||
for j,ji in enumerate(jit_cache):
|
||||
op_estimate += ji.prg.op_estimate
|
||||
mem_estimate += ji.prg.mem_estimate
|
||||
lds_estimate += ji.prg.lds_estimate
|
||||
if isinstance(ji.prg, CompiledRunner):
|
||||
if ji.prg.p.vars: self.jc_idx_with_updatable_var_vals.append(j)
|
||||
if (ji.prg.p.global_size and not all_int(ji.prg.p.global_size)) or (ji.prg.p.local_size and not all_int(ji.prg.p.local_size)):
|
||||
self.jc_idx_with_updatable_launch_dims.append(j)
|
||||
self.vars = sorted(var_vals.keys(), key=lambda v: v.expr)
|
||||
if ji.prg.p.vars: self.var_vals_replace[j] = [self.vars.index(v) for v in ji.prg.p.vars]
|
||||
|
||||
global_dim_idx, local_dim_idx = find_symbolic_dim(ji.prg.p.global_size), find_symbolic_dim(ji.prg.p.local_size)
|
||||
if global_dim_idx is not None or local_dim_idx is not None: self.launch_dims_replace[j] = (global_dim_idx, local_dim_idx)
|
||||
|
||||
super().__init__(colored(f"<batched {len(self.jit_cache)}>", "cyan"), jit_cache[0].prg.dname.split(":")[0],
|
||||
op_estimate, mem_estimate, lds_estimate)
|
||||
|
||||
def updated_vars(self, var_vals):
|
||||
vals = [var_vals[v] for v in self.vars]
|
||||
for j, vidxs in self.var_vals_replace.items():
|
||||
for i, v in enumerate(vidxs): yield j, i, vals[v]
|
||||
|
||||
def updated_launch_dims(self, var_vals):
|
||||
dims = [tuple(sym_infer(s, var_vals) for s in dim) for dim in self.symbolic_dims]
|
||||
for j, (gl, lc) in self.launch_dims_replace.items(): yield j, (dims[gl] if gl is not None else None), (dims[lc] if lc is not None else None)
|
||||
|
||||
class MultiGraphRunner(GraphRunner): # pylint: disable=abstract-method
|
||||
def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
|
||||
self.w_dependency_map: Dict[Any, Any] = {}
|
||||
|
||||
@@ -33,7 +33,7 @@ class CUDAGraph(MultiGraphRunner):
|
||||
kern_params = cuda.CUDA_KERNEL_NODE_PARAMS(ji.prg.clprg.prg, *global_size, *local_size, 0, None, vargs)
|
||||
check(cuda.cuGraphAddKernelNode(ctypes.byref(new_node), self.graph, c_deps, len(deps), ctypes.byref(kern_params)))
|
||||
|
||||
if j in self.jc_idx_with_updatable_launch_dims or j in self.jc_idx_with_updatable_var_vals or j in self.jc_idx_with_updatable_rawbufs:
|
||||
if j in self.launch_dims_replace or j in self.var_vals_replace or j in self.jc_idx_with_updatable_rawbufs:
|
||||
self.updatable_nodes[j] = (new_node, kern_params, c_args, False)
|
||||
elif isinstance(ji.prg, BufferXfer):
|
||||
dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
|
||||
@@ -58,13 +58,13 @@ class CUDAGraph(MultiGraphRunner):
|
||||
elif i == 1: self.updatable_nodes[j][1].srcDevice = input_rawbuffers[input_idx]._buf
|
||||
|
||||
# Update var_vals in the c_args struct.
|
||||
for j in self.jc_idx_with_updatable_var_vals:
|
||||
for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars):
|
||||
setattr(self.updatable_nodes[j][2], f'v{i}', var_vals[v])
|
||||
for j, i, v in self.updated_vars(var_vals): setattr(self.updatable_nodes[j][2], f'v{i}', v)
|
||||
|
||||
# Update launch dims in the kern_params struct.
|
||||
for j in self.jc_idx_with_updatable_launch_dims:
|
||||
self.set_kernel_node_launch_dims(self.updatable_nodes[j][1], *cast(CompiledRunner, self.jit_cache[j].prg).p.launch_dims(var_vals))
|
||||
for j, global_dims, local_dims in self.updated_launch_dims(var_vals):
|
||||
prg = cast(CompiledRunner, self.jit_cache[j].prg)
|
||||
node, global_size, local_size = self.updatable_nodes[j][1], global_dims or prg.p.global_size, local_dims or prg.p.local_size
|
||||
node.blockDimX, node.blockDimY, node.blockDimZ, node.gridDimX, node.gridDimY, node.gridDimZ = *local_size, *global_size # type: ignore[misc]
|
||||
|
||||
# Update graph nodes with the updated structs.
|
||||
for node, c_node_params, c_args, is_copy in self.updatable_nodes.values():
|
||||
@@ -76,6 +76,3 @@ class CUDAGraph(MultiGraphRunner):
|
||||
def __del__(self):
|
||||
if hasattr(self, 'graph'): check(cuda.cuGraphDestroy(self.graph))
|
||||
if hasattr(self, 'instance'): check(cuda.cuGraphExecDestroy(self.instance))
|
||||
|
||||
def set_kernel_node_launch_dims(self, node, global_size: Tuple[int, int, int], local_size: Tuple[int, int, int]):
|
||||
node.blockDimX, node.blockDimY, node.blockDimZ, node.gridDimX, node.gridDimY, node.gridDimZ = *local_size, *global_size
|
||||
|
||||
@@ -137,12 +137,12 @@ class HCQGraph(MultiGraphRunner):
|
||||
else: self.op_cmd_idx[j][0].update_copy(self.op_cmd_idx[j][1], **{('dest' if i == 0 else 'src'): input_rawbuffers[input_idx]._buf.va_addr})
|
||||
|
||||
# Update var_vals
|
||||
for j in self.jc_idx_with_updatable_var_vals:
|
||||
for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars): self.ji_args_vars[j][i] = var_vals[v]
|
||||
for j, i, v in self.updated_vars(var_vals): self.ji_args_vars[j][i] = v
|
||||
|
||||
for j in self.jc_idx_with_updatable_launch_dims:
|
||||
# Update launch dims
|
||||
for j, global_dims, local_dims in self.updated_launch_dims(var_vals):
|
||||
queue, cmd_ptr = self.op_cmd_idx[j]
|
||||
queue.update_exec(cmd_ptr, *cast(CompiledRunner, self.jit_cache[j].prg).p.launch_dims(var_vals))
|
||||
queue.update_exec(cmd_ptr, global_dims, local_dims)
|
||||
|
||||
for dev in self.devices:
|
||||
comp_queue, copy_queue, need_sig_upd = self.comp_queues[dev], self.copy_queues[dev], dev.timeline_signal != self.last_timeline[dev][0]
|
||||
|
||||
@@ -39,9 +39,9 @@ class MetalGraph(GraphRunner):
|
||||
icb_command.setKernelBuffer_offset_atIndex_(b._buf.buf, b._buf.offset, i)
|
||||
all_resources.append(b._buf.buf)
|
||||
for i,v in enumerate(prg.p.vars): icb_command.setKernelBuffer_offset_atIndex_(self.int_buf.buf, self.vars.index(v)*4, len(ji.bufs)+i)
|
||||
if j not in self.jc_idx_with_updatable_launch_dims:
|
||||
global_size, local_size = prg.p.launch_dims(var_vals)
|
||||
icb_command.concurrentDispatchThreadgroups_threadsPerThreadgroup_(Metal.MTLSize(*global_size), Metal.MTLSize(*local_size))
|
||||
|
||||
global_size, local_size = prg.p.launch_dims(var_vals)
|
||||
icb_command.concurrentDispatchThreadgroups_threadsPerThreadgroup_(Metal.MTLSize(*global_size), Metal.MTLSize(*local_size))
|
||||
icb_command.setBarrier()
|
||||
|
||||
self.all_resources = dedup(all_resources)
|
||||
@@ -55,8 +55,10 @@ class MetalGraph(GraphRunner):
|
||||
for (j,i),input_idx in self.input_replace.items():
|
||||
self.icb.indirectComputeCommandAtIndex_(j).setKernelBuffer_offset_atIndex_(input_rawbuffers[input_idx]._buf.buf,
|
||||
input_rawbuffers[input_idx]._buf.offset, i)
|
||||
for j in self.jc_idx_with_updatable_launch_dims:
|
||||
global_size, local_size = cast(CompiledRunner, self.jit_cache[j].prg).p.launch_dims(var_vals)
|
||||
|
||||
for j, global_dims, local_dims in self.updated_launch_dims(var_vals):
|
||||
prg = cast(CompiledRunner, self.jit_cache[j].prg)
|
||||
global_size, local_size = global_dims or prg.p.global_size, local_dims or prg.p.local_size
|
||||
self.icb.indirectComputeCommandAtIndex_(j).concurrentDispatchThreadgroups_threadsPerThreadgroup_(Metal.MTLSize(*global_size),
|
||||
Metal.MTLSize(*local_size))
|
||||
for j, var in enumerate(self.vars): self.int_buf_view[j] = var_vals[var]
|
||||
|
||||
Reference in New Issue
Block a user