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

hcq2: fix debug 2 info (#17491)

Browse Source 
* hcq2: fix debug 2 info

* x

* x

* x
This commit is contained in:
nimlgen
2026-08-12 00:25:44 +03:00
committed by GitHub
parent e11df72e0f
commit 2b5018e86a
2 changed files with 25 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/engine/realize.py
+18 -15
View File
@@ -5,7 +5,7 @@ from dataclasses import dataclass, replace, field
from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite
from tinygrad.device import Device, Buffer, MultiBuffer
from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
from tinygrad.renderer import Estimates
from tinygrad.codegen import to_program
from tinygrad.codegen.opt.postrange import args_from_ast
@@ -210,27 +210,30 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
return t[0]
def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
if (inputs:=call.arg.aux.inputs) is not None:
if (info:=call.arg.aux).inputs is not None:
bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
table = call.src[1+inputs].buffer
table = call.src[1+info.inputs].buffer
for j,dev in enumerate(call.arg.aux.device):
addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q')
wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected")
mv[:len(addrs)] = addrs
exec_kernel(replace(ctx, update_stats=False), call, ast)
exec_kernel(replace(ctx, update_stats=DEBUG>=3), call, ast)
tms:list[float|None] = []
for e in (aux:=call.arg.aux).prof: cast(Any, Device[e.device]).prof_ents[e.st_id] = e
for d in [cast(Any, Device[x]) for x in aux.device]:
with track_stats(ctx, call, d.device, [], ctx.var_vals) as et:
if ctx.wait:
d.synchronize(timeout=ctx.timeout)
ts = [d.signal(i)._buf.cpu_view().view(fmt='Q')[0] for e in aux.prof if e.device == d.device for i in (e.st_id, e.en_id)]
if ts: et[0] = float(max(ts)-min(ts))/d.timestamp_divider/1e6
tms += et
return tms[0]
tms = []
for devices,name,estimates,prof in info.kernels:
for device in devices:
d, tm = cast(Any, Device[device]), None
if prof:
d.prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
if ctx.wait:
d.synchronize(timeout=ctx.timeout)
st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
with track_stats(ctx, call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates))), d.device, [], ctx.var_vals) as et:
et[0] = tm
return max(tms) if tms else None
# flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
pm_flatten_linear = PatternMatcher([
@@ -276,7 +279,7 @@ def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:li
if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE) if profile is None else profile)
if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear
tinygrad/runtime/support/hcq2.py
+7 -7
View File
@@ -33,7 +33,7 @@ class HCQInfo:
input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
inputs:int|None = None
prof:tuple[ProfileGraphEntry, ...] = () # st_id/en_id are timestamp signal slots until collect
kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...]], ...] = ()
def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
@@ -199,10 +199,10 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
signal_tags |= cur_signal_tags
# build fences and finalizers
fences, finalizers, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
fences, fins, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
signal_tags |= finalizer_signal_tags
src, prof = [], []
src, kerns = [], []
for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
# first queue use, sync prior device work with the device timeline
if batch_info.index((devices, queue)) == tag:
@@ -212,18 +212,18 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
# and make hcq call
name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call))
ts_ids = [next(UOp.unique_num) for _ in range(2)] if profile else []
prof += [ProfileGraphEntry(d, name, *ts_ids) for d in devices if ts_ids]
kerns.append((devices, name, info.estimates, tuple(ts_ids)))
ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_signal(devices, s),)) for s in ts_ids]
q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:]
# signal the queue if someone waits for us
if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
src.append(make_call(name, make_submit(*q, devs=devices, queue=queue).sink(), info))
src.append(make_call(f"submit {name}", make_submit(*q, devs=devices, queue=queue).sink(), info))
# append batch timestamps to finalizers
finalizers = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, prof=tuple(e for e in prof if e.device in a.device)))) for f in finalizers]
return fences + src + finalizers
fins = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, kernels=tuple(x for x in kerns if set(x[0]) & set(a.device))))) for f in fins]
return fences + src + fins
def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
srcs:list[UOp] = []