diff --git a/extra/hcq2/hcq2.py b/extra/hcq2/hcq2.py index b5eb00dab4..58bddb7b0a 100644 --- a/extra/hcq2/hcq2.py +++ b/extra/hcq2/hcq2.py @@ -49,7 +49,7 @@ def make_getaddr(u, device=None): return UOp(Ops.GETADDR, dtypes.uint64, src=(u,), arg=device or to_tuple(u.device)[0]) def make_ins(op, *srcs): - return UOp(Ops.INS, dtypes.void, tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs), op) + return UOp(Ops.INS, arg=op, src=tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs)) def make_placeholder(devs, size:int, dtype, name=None, unique=True) -> UOp: return UOp.param(next(UOp.unique_num) if unique else 0, dtype, shape=(size,), device=devs).rtag(name or "temp") @@ -133,7 +133,7 @@ def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, for (ddevs, dqueue, dtag), lanes in deps.items(): sig = make_mstack([make_signal(d if dl is None else ddevs[dl], queue=dqueue, sentinel=dl is None) for dl, d in zip(lanes, devices)]) val = make_mstack([make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)]) - waits.append((sig.index(zero:=UOp.const(dtypes.int, 0)).load() >= val.index(zero) + dtag).wait()) + waits.append(UOp(Ops.INS, arg="wait", src=(sig, val.index(UOp.const(dtypes.int, 0)) + dtag))) return waits, {dtag for _, _, dtag in deps} def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]], @@ -154,7 +154,8 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t waited |= cur_waited # wait the syncs, store the device epoch; value bumps are a separate call: no lane may bump until every lane has patched its waits - submit = make_submit(*waits, make_signal(devs).store((tl:=make_signal_value(devs)).index(zero)), devs=devs, queue="COMPUTE:0") + store = UOp(Ops.INS, arg="store", src=(make_signal(devs), (tl:=make_signal_value(devs)).index(zero))) + submit = make_submit(*waits, store, devs=devs, queue="COMPUTE:0") upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), n) for qn in dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)])] bump = UOp.barrier(*[s.index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd]) finalizers += [UOp.custom_function("hcq", b.sink()).call(aux=HCQInfo("hcq_finalizer", Estimates(), devs, "COMPUTE:0")) for b in (submit, bump)] @@ -181,15 +182,15 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]: for tag, ((call, _), (devices, queue), cmds) in enumerate(zip(batch, batch_info, call_waits)): # first queue use, sync prior device work with main signal if batch_info.index((devices, queue)) == tag: - epoch = (make_signal(devices).index(0).load() >= make_signal_value(devices).index(0) - 1).wait() - cmds = [UOp(Ops.BARRIER), epoch] + cmds - - # signal queue timeline if someone waits for us - store = make_signal(devices, queue=queue).store(make_signal_value(devices, queue=queue).index(0) + tag) if tag in waited else None + epoch = UOp(Ops.INS, arg="wait", src=(make_signal(devices), make_signal_value(devices).index(0) - 1)) + cmds = [UOp(Ops.INS, arg="barrier", src=()), epoch] + cmds # and make hcq call info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), devices, queue) - cmds = [*cmds, call.replace(arg=replace(call.arg, aux=info))] + ([store] if store is not None else []) + cmds = [*cmds, call.replace(arg=replace(call.arg, aux=info))] + + # signal queue timeline if someone waits for us + if tag in waited: cmds += [UOp(Ops.INS, arg="store", src=(make_signal(devices, queue), make_signal_value(devices, queue).index(0) + tag))] src.append(UOp.custom_function("hcq", make_submit(*cmds, devs=devices, queue=queue).sink()).call(name="hcq", aux=info)) return src + finalizers diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index 184e831e09..83ae57c979 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -90,7 +90,7 @@ def memory_barrier(ctx): reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff), acquire_mem(ctx))) -def pm4_wait(ctx, x, y): return wait_reg_mem(ctx, y, mem=make_getaddr(x.buf_uop, ctx.devs)) +def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=make_getaddr(dst, ctx.devs)) def pm4_barrier(ctx): return memory_barrier(ctx) @@ -138,10 +138,10 @@ def pm4_program(ctx, call, prg): pm_pm4_opsel = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program), - (UPat(Ops.WAIT, src=(UPat.var("x") >= UPat.var("y"),)), pm4_wait), - (UPat(Ops.BARRIER), pm4_barrier), - (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp), - (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store), + (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), pm4_wait), + (UPat(Ops.INS, arg="barrier"), pm4_barrier), + (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp), + (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store), ]) def pm4_submit(cmdbuf, devs): @@ -184,10 +184,10 @@ def sdma_copy(ctx, call): ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz - off, ctx.max_copy_size) - 1), 0, *data64_le(src_addr + off), *data64_le(dst_addr + off)) for off in range(0, sz, ctx.max_copy_size)])) -def sdma_wait(ctx, x, y): +def sdma_wait(ctx, dst, val): op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \ | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1) - return make_ins(SDMAOps.POLL_REGMEM, op, *data64_le(make_getaddr(x.buf_uop, ctx.devs)), y, 0xffffffff, + return make_ins(SDMAOps.POLL_REGMEM, op, *data64_le(make_getaddr(dst, ctx.devs)), val, 0xffffffff, ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff)) def sdma_store(ctx, dst, val): @@ -202,10 +202,10 @@ def sdma_timestamp(ctx, dst): pm_sdma_opsel = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy), - (UPat(Ops.BARRIER), lambda: UOp(Ops.NOOP, dtypes.void, ())), - (UPat(Ops.WAIT, src=(UPat.var("x") >= UPat.var("y"),)), sdma_wait), - (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), sdma_timestamp), - (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store), + (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())), + (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), sdma_wait), + (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), sdma_timestamp), + (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store), ]) def sdma_submit(cmdbuf, devs):