diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 026e35977c..fec53400be 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1434,9 +1434,9 @@ def train_llama3(): load_state_dict(scheduler, safe_load(fn), realize=False) fp8_amax = [t for ts in model._fp8_amax.values() for t in ts] - fp8_next_amax = [t for ts in model._fp8_next_amax.values() for t in ts] if hasattr(model, "_fp8_next_amax") else [] - fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] if hasattr(model, "_fp8_grad_amax") else [] - fp8_next_grad_amax = [t for ts in model._fp8_next_grad_amax.values() for t in ts] if hasattr(model, "_fp8_next_grad_amax") else [] + fp8_next_amax = [t for ts in model._fp8_next_amax.values() for t in ts] + fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] + fp8_next_grad_amax = [t for ts in model._fp8_next_grad_amax.values() for t in ts] fp8_inv_scales = list(model._fp8_inv_scale.values()) + list(model._fp8_next_inv_scale.values()) from tinygrad.nn.state import get_state_dict @@ -1462,8 +1462,7 @@ def train_llama3(): @TinyJit def minibatch(tokens:Tensor): - for nxt in fp8_next_amax: nxt.assign(0) - for nxt in fp8_next_grad_amax: nxt.assign(0) + model.reset_amax() if is_dp: tokens = tokens.to(None).shard(device, 0) if is_mp: tokens = tokens.shard(device) if not is_sharding: tokens = tokens.to(None) @@ -1487,8 +1486,7 @@ def train_llama3(): scheduler.step() for g in grads: g.assign(0) - for cur, nxt in zip(fp8_amax, fp8_next_amax): cur.assign(nxt) - for cur, nxt in zip(fp8_grad_amax, fp8_next_grad_amax): cur.assign(nxt) + model.update_amax() lr_cpu = optim.lr.float().to("CPU") grad_norm_cpu = grad_norm.float().to("CPU") diff --git a/examples/mlperf/models/flat_llama.py b/examples/mlperf/models/flat_llama.py index 8cbc0e9223..d237823687 100644 --- a/examples/mlperf/models/flat_llama.py +++ b/examples/mlperf/models/flat_llama.py @@ -83,8 +83,8 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca return out, x_fp8 return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8 -def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor, - next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor): +def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None, + next_amax_x:Tensor|None, grad_amax_state:Tensor|None, next_grad_amax_state:Tensor|None): if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4: from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8 x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x) @@ -96,8 +96,8 @@ def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, ep next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x) return out, x_normed, rrms, ret -def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor, - next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None): +def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None, + next_amax_x:Tensor|None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None): if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4: from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8 x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x) @@ -111,9 +111,9 @@ def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w return out, h, x_normed, rrms, ret def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor, - amax_x2:Tensor, next_amax_x2:Tensor, - grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor, - grad_amax_xout:Tensor, next_grad_amax_xout:Tensor): + amax_x2:Tensor|None, next_amax_x2:Tensor|None, + grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None, + grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None): if FUSED_SILU_W13 and not MXFP4: from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13 x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13, @@ -164,14 +164,15 @@ class FlatTransformer: self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False) def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False) + n_amax = 0 if MXFP4 else n_layers names = ["xqkv", "xo", "x2"] names += ["x1", "x3"] if SPLIT_W13 else ["x13"] - self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names} - self._fp8_next_amax = {name: [_amax() for _ in range(n_layers)] for name in names} + self._fp8_amax = {name: [_amax() for _ in range(n_amax)] for name in names} + self._fp8_next_amax = {name: [_amax() for _ in range(n_amax)] for name in names} grad_names = ["xqkv", "xo", "xout"] grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"] - self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names} - self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names} + self._fp8_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names} + self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names} w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)] w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)] self._fp8_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales} @@ -195,9 +196,10 @@ class FlatTransformer: return (w * scale_b).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor, - amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor, - next_amax_xqkv:Tensor, next_amax_xo:Tensor, - grad_amax_xqkv:Tensor, grad_amax_xo:Tensor, next_grad_amax_xqkv:Tensor, next_grad_amax_xo:Tensor): + amax_xqkv:Tensor|None, amax_xo:Tensor|None, s_qkv:Tensor, s_o:Tensor, + next_amax_xqkv:Tensor|None, next_amax_xo:Tensor|None, + grad_amax_xqkv:Tensor|None, grad_amax_xo:Tensor|None, + next_grad_amax_xqkv:Tensor|None, next_grad_amax_xo:Tensor|None): bsz, seqlen, _ = x.shape saves = [] @@ -319,28 +321,33 @@ class FlatTransformer: for i in range(len(amax_dict[name])): amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False) + def reset_amax(self): + for st in (self._fp8_next_amax, self._fp8_next_grad_amax): + for ts in st.values(): + for t in ts: t.assign(0) + + def update_amax(self): + for cur, nxt in ((self._fp8_amax, self._fp8_next_amax), (self._fp8_grad_amax, self._fp8_next_grad_amax)): + for name in cur: + for c, n in zip(cur[name], nxt[name]): c.assign(n) + def __call__(self, tokens:Tensor, save:bool=True): h = self.tok_embeddings(tokens) freqs_cis = self.freqs_cis.cast(h.dtype) if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :] a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale + def amax_kwargs(i:int, act_names:tuple[str, ...], grad_names:tuple[str, ...]) -> dict[str, Tensor|None]: + specs = (("amax_", a, act_names), ("next_amax_", na, act_names), ("grad_amax_", ga, grad_names), ("next_grad_amax_", nga, grad_names)) + if MXFP4: return dict.fromkeys(f"{prefix}{name}" for prefix, _, names in specs for name in names) + return {f"{prefix}{name}":val[name][i] for prefix, val, names in specs for name in names} for i in range(self.n_layers): - attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i], - amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i], s_qkv=s["wqkv"][i], s_o=s["wo"][i], - next_amax_xqkv=na["xqkv"][i], next_amax_xo=na["xo"][i], - grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i], - next_grad_amax_xqkv=nga["xqkv"][i], next_grad_amax_xo=nga["xo"][i]) - ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i], - amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i], - next_amax_x2=na["x2"][i]) + attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i], s_qkv=s["wqkv"][i], s_o=s["wo"][i], + **amax_kwargs(i, ("xqkv", "xo"), ("xqkv", "xo"))) + ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i], s_2=s["w2"][i], **amax_kwargs(i, ("x2",), ("xout",))) if SPLIT_W13: - ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i], - next_amax_x1=na["x1"][i], next_amax_x3=na["x3"][i], - s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i], - next_grad_amax_xw1=nga["xw1"][i], next_grad_amax_xw3=nga["xw3"][i]) + ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], s_1=s["w1"][i], s_3=s["w3"][i], **amax_kwargs(i, ("x1", "x3"), ("xw1", "xw3"))) else: - ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i], - next_grad_amax_xw13=nga["xw13"][i], next_amax_x13=na["x13"][i]) + ffn_kwargs.update(w13=self.w13[i], s_13=s["w13"][i], **amax_kwargs(i, ("x13",), ("xw13",))) h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save) logits = matmul(self.norm(h), self.output[0], fp8=False)[0] @@ -424,9 +431,7 @@ if __name__ == "__main__": @TinyJit def fwd_bwd(tokens:Tensor): with Timing("python forward: "): - for amax_dict in (model._fp8_next_amax, model._fp8_next_grad_amax): - for ts in amax_dict.values(): - for nxt in ts: nxt.assign(0) + model.reset_amax() logits = model(tokens[:, :-1], save=llama_size=="8B") loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:]) with Timing("python backward: "): diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index 38a76c257a..788d40f495 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -4,7 +4,7 @@ import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, co assert sys.platform != 'win32' from dataclasses import dataclass from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf -from tinygrad.runtime.support.hcq2 import make_binary_patch, make_patches +from tinygrad.runtime.support.hcq2 import make_binary_patch from tinygrad.uop.ops import sint, UOp from tinygrad.device import Compiled, BufferSpec, Buffer, Device from tinygrad.dtype import dtypes @@ -152,19 +152,12 @@ def pm4_submit(ctx, lin): ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}") for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value))) - # two tail dwords coordinate safe IB reuse: GPU completions and host submits - size_dw = sum(len(ins.src) for ins in lin.src) + len(release_mem(ctx, 0, 0).src) + # the host fence at the start of the batch guarantees the ib is free to reuse + size_dw = sum(len(ins.src) for ins in lin.src) assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet" - ib = UOp.placeholder((size_dw + 2,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf") - done_idx, submit_idx = UOp.const(size_dw + 0, dtypes.int), UOp.const(size_dw + 1, dtypes.int) - init_counters = make_patches(ib, [((size_dw + i) * 4, UOp.const(0, dtypes.uint32)) for i in range(2)]).rtag("link") - submitted = (counter:=ib.after(init_counters).index(submit_idx)).load() - completed = ib.after(loop:=UOp.loop(0)).index(done_idx).load() - ib_free = completed.end(loop, completed != submitted) - - bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(size_dw)), 2), (submitted + 1).cast(dtypes.uint64)) - cmdbuf = make_cmdbuf(lin.replace(src=lin.src + (bump_fence,)), devs, buf=ib, dep=ib_free) + ib = UOp.placeholder((size_dw,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf") + cmdbuf = make_cmdbuf(lin, devs, buf=ib) # the ring itself only carries a packet pointing at the ib, wrapping the ring put = put_ptr.index(zero:=UOp.const(0, dtypes.int)) @@ -174,7 +167,7 @@ def pm4_submit(ctx, lin): # advance the put/write pointers past the packet bump_put_ptr = put_ptr.index(zero).store(put + len(pkt)) bump_wptr = wptr.index(zero).store(put + len(pkt)) - flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr, counter.store(submitted + 1)) + flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr) return doorbell.after(flush).index(zero).store(put + len(pkt)) pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)]) @@ -518,8 +511,7 @@ class PCIIface(PCIIfaceBase): cq = d.compute_queue for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0 d.iface.dev_impl.gfx.setup_ring(*cq.params) - d.timeline_signal('COMPUTE:0')._buf.cpu_view().mv.cast('Q')[0] = \ - d.timeline_value('COMPUTE:0').as_memoryview(force_zero_copy=True).cast('Q')[0] - 1 + d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = d.signal('value', 1).as_memoryview(force_zero_copy=True).cast('Q')[0] - 1 def sleep(self, timeout): if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))): @@ -639,9 +631,6 @@ class AMDDevice(HCQ2Compiled): qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}" self.pm_bufferize = PatternMatcher([ (UPat(Ops.PARAM, tag=f"{qname}_{name}"), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"] - ] + [ - (UPat(Ops.PARAM, tag=f"{qname}_timeline_signal"), lambda ctx, q=qname: ctx[0].timeline_signal(q)), - (UPat(Ops.PARAM, tag=f"{qname}_timeline_value"), lambda ctx, q=qname: ctx[0].timeline_value(q)), ]) + self.pm_bufferize return queue diff --git a/test/device/test_hcq.py b/test/device/test_hcq.py index 369adddf8d..b8390d0b86 100644 --- a/test/device/test_hcq.py +++ b/test/device/test_hcq.py @@ -330,6 +330,7 @@ class TestHCQ(unittest.TestCase): # Test profile api def test_speed_exec_time(self): sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal() + st = time.perf_counter() TestHCQ.d0.hw_compute_queue_t().timestamp(sig_st) \ .exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \ .timestamp(sig_en) \ @@ -337,11 +338,13 @@ class TestHCQ(unittest.TestCase): TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value) TestHCQ.d0.timeline_value += 1 + host_us = (time.perf_counter() - st) * 1e6 et = float(sig_en.timestamp - sig_st.timestamp) print(f"exec kernel time: {et:.2f} us") - assert 0.1 <= et <= (3000000 if MOCKGPU or Device.DEFAULT in {"CPU"} else 100) + # emulated devices are only bounded by the host window around submit+wait + assert 0.1 <= et <= (host_us if MOCKGPU or Device.DEFAULT in {"CPU"} else 100) def test_speed_copy_bandwidth(self): if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue") @@ -352,6 +355,7 @@ class TestHCQ(unittest.TestCase): b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate() sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal() + st = time.perf_counter() TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \ .copy(a._buf, b._buf, SZ) \ .timestamp(sig_en) \ @@ -359,13 +363,14 @@ class TestHCQ(unittest.TestCase): TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value) TestHCQ.d0.timeline_value += 1 + host_ms = (time.perf_counter() - st) * 1e3 - et = float(sig_en.timestamp - sig_st.timestamp) - et_ms = et / 1e3 + et_ms = float(sig_en.timestamp - sig_st.timestamp) / 1e3 + assert 0 < et_ms <= host_ms # timestamps are in us and cover only the copy gb_s = ((SZ / 1e9) / et_ms) * 1e3 print(f"same device copy: {et_ms:.2f} ms, {gb_s:.2f} GB/s") - assert (0.2 if MOCKGPU else 10) <= gb_s <= 1000 + assert (0 if MOCKGPU else 10) <= gb_s <= 1000 def test_speed_cross_device_copy_bandwidth(self): if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue") @@ -379,6 +384,7 @@ class TestHCQ(unittest.TestCase): TestHCQ.d0.allocator._map(b._buf) sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal() + st = time.perf_counter() TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \ .copy(a._buf, b._buf, SZ) \ .timestamp(sig_en) \ @@ -386,13 +392,14 @@ class TestHCQ(unittest.TestCase): TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value) TestHCQ.d0.timeline_value += 1 + host_ms = (time.perf_counter() - st) * 1e3 - et = float(sig_en.timestamp - sig_st.timestamp) - et_ms = et / 1e3 + et_ms = float(sig_en.timestamp - sig_st.timestamp) / 1e3 + assert 0 < et_ms <= host_ms # timestamps are in us and cover only the copy gb_s = ((SZ / 1e9) / et_ms) * 1e3 print(f"cross device copy: {et_ms:.2f} ms, {gb_s:.2f} GB/s") - assert (0.2 if MOCKGPU else 2) <= gb_s <= 100 + assert (0 if MOCKGPU else 2) <= gb_s <= 100 def test_timeline_signal_rollover(self): for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]: diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 40d6636e56..97074388ff 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -71,7 +71,7 @@ def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp: # parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index linear = linear.substitute({u: UOp.param(i, u.dtype, u.shape, u.device) for i,u in enumerate(input_uops)}, walk=True) linear = memory_plan_rewrite(linear, held_bufs) - linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value), jit=True) + linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value)) if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value) if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View graphed linear") return linear @@ -169,7 +169,7 @@ class CapturedJit(Generic[ReturnType]): expected_input_info: list[tuple[UOp, tuple[Variable, ...], DType, str]] # (view, variables, dtype, device) per input @functools.cached_property - def linear(self) -> UOp: return link_linear(self._linear, jit=True) + def linear(self) -> UOp: return link_linear(self._linear) def __reduce__(self): return self.__class__, (self.ret, self._linear, self.expected_names, self.expected_input_info) diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 91b9e331a6..38d5c2831d 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -265,18 +265,18 @@ pm_exec = PatternMatcher([ if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link # noqa: E402 # down here, hcq2 imports the helpers above -def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, jit=False) -> UOp: +def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None) -> UOp: 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, jit=jit) + if getenv("HCQ2"): linear = hcq_compile(linear, input_uops) return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True) -def link_linear(linear:UOp, jit=False, cache=True) -> UOp: return hcq_link(linear, jit=jit, cache=cache) if getenv("HCQ2") else linear +def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False): inputs = list(input_uops) - if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs, jit=False)) + if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs)) ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2) for call in linear.src: pm_exec.rewrite(call, ctx) @@ -287,4 +287,5 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None from tinygrad.tensor import Tensor with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False) ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False) - return pm_exec.rewrite(link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0), cache=ctx.cache).src[0], ctx) + linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0), cache=ctx.cache) + return max(pm_exec.rewrite(c, ctx) or 0.0 for c in linear.src) diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index af0bf2b406..c74363e659 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -164,13 +164,13 @@ class CPUDevice(HCQCompiled): (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_host_queue)]) pm_bufferize = PatternMatcher([ - (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].timeline("sentinel", (1 << 64) - 1)), - (UPat(Ops.PARAM, tag="COMPUTE:0_timeline_signal"), lambda ctx: ctx[0].timeline("signal", 0)), - (UPat(Ops.PARAM, tag="COMPUTE:0_timeline_value"), lambda ctx: ctx[0].timeline("value", 1)), + (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)), + (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].signal("timeline")), + (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].signal("value", 1)), ]) @functools.cache - def timeline(self, tag:str, init_value:int) -> Buffer: + def signal(self, name:str, init_value:int=0) -> Buffer: (buf:=Buffer(self.device, 1, dtypes.uint64, preallocate=True)).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value return buf diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 008167ae3f..e01dab947f 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -41,28 +41,33 @@ def unwrap_mstack(u): if u.op is Ops.MSTACK: return tuple(x for s in u.src for x in unwrap_mstack(s)) return unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,) -def make_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> UOp: - offsets = UOp(Ops.STACK, dtypes.int, tuple(UOp.const(off // buf.dtype.itemsize, dtypes.int) for off,_ in patches)) - values = UOp(Ops.STACK, buf.dtype, tuple(val.cast(buf.dtype) for _,val in patches)) - return buf.index(offsets).store(values) +def is_value_known_at_link(val:UOp) -> bool: + runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)] + addressed_bufs = [b for g in val.toposort() if g.op is Ops.GETADDR for b in unwrap_mstack(g.buf_uop)] + + # addr of input params is not known at link time + return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs) + +def make_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> tuple[UOp, ...]: + return tuple(buf.index(UOp(Ops.STACK, dtypes.int, tuple(UOp.const(off // buf.dtype.itemsize, dtypes.int) for off,_ in ps))) + .store(UOp(Ops.STACK, buf.dtype, tuple(val.cast(buf.dtype) for _,val in ps))).rtag(tag) + for ps, tag in zip(partition(patches, lambda p: is_value_known_at_link(p[1])), ("link", None)) if ps) def make_binary_patch(buf:UOp, blob:bytes) -> UOp: data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype) r = UOp.range(len(blob) // buf.dtype.itemsize, 0, dtype=dtypes.int, src=(buf, data)) - return buf.index(r).store(data.index(r).load()).end(r) + return buf.index(r).store(data.index(r).load()).end(r).rtag("link") -def make_cmdbuf(lin, devs, buf:UOp|None=None, dep:UOp|None=None): +def make_cmdbuf(lin, devs, buf:UOp|None=None): blob, patches = bytearray(), [] for s in (s for ins in lin.src for s in ins.src): if s.op is not Ops.CONST: patches.append((len(blob), s)) blob.extend(struct.pack(f'<{s.dtype.fmt}', s.val if s.op is Ops.CONST else 0x0)) cmdbuf = buf if buf is not None else UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf") - writable = cmdbuf.after(dep) if dep is not None else cmdbuf - return cmdbuf.after(make_binary_patch(writable, bytes(blob)), *((make_patches(writable, patches),) if patches else ())) + return cmdbuf.after(make_binary_patch(cmdbuf, bytes(blob)), *make_patches(cmdbuf, patches)) -def make_signal(devs, queue="COMPUTE:0", sentinel=False): - return UOp.placeholder((1,), dtypes.uint64, 0, device=devs, volatile=True).rtag("sentinel_signal" if sentinel else f"{queue}_timeline_signal") -def make_signal_value(devs, queue="COMPUTE:0"): return UOp.placeholder((1,), dtypes.uint64, 0, device=devs).rtag(f"{queue}_timeline_value") +def make_signal(devs, slot:int=0, tag:str="signal") -> UOp: + return UOp.placeholder((1,), dtypes.uint64, slot, device=devs, volatile=True).rtag(tag) def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp: return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, src=tuple(cmds), arg=(to_tuple(devs), queue))) @@ -72,7 +77,7 @@ def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp: data, info = prg.arg buf = UOp.placeholder((data.kernargs_alloc_size // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("kernargs") words = [w for gi in info.globals for w in data64_le(get_call_arg_uops(call)[gi].getaddr(devs))] + list(info.vars) - return buf.after(*((make_patches(buf, [(i * 4, w) for i, w in enumerate(words)]),) if words else ())) + return buf.after(*make_patches(buf, [(i * 4, w) for i, w in enumerate(words)])) # ***************** # 0.1. prep: replace buffers with params @@ -115,7 +120,7 @@ def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tu dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, written, (key, lane))] return dep_lanes -def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]: +def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]: # opt1: same-queue ops are fifo-ordered if devices[0].split(":")[0] in {"AMD", "QCOM"} or queue.startswith("COPY"): dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep[0][dlane], dep[1]) != (devices[lane], queue)] @@ -127,70 +132,81 @@ def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, waits = [] for (ddevs, dqueue, dtag), lanes in deps.items(): - sig = UOp.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 = UOp.mstack(*[make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)]) - waits.append(UOp(Ops.INS, arg="wait", src=(sig, val.index(UOp.const(0, dtypes.int)) + dtag))) + sig = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue]) + for dl, d in zip(lanes, devices)]) + waits.append(UOp(Ops.INS, arg="wait", src=(sig, UOp.const(dtag + 1, dtypes.uint64)))) return waits, {dtag for _, _, dtag in deps} +def make_fence(timeline:UOp, prev:UOp, sigs:list[UOp]) -> UOp: + free = (cur:=timeline.after(loop:=UOp.loop(0)).index(0).load()).end(loop, cur < prev.index(0).load()) + return UOp.sink(*[s.after(free).index(0).store(0) for s in sigs]) + +def _hcq_call(devs, name:str, body:UOp) -> UOp: return UOp.custom_function("hcq", body).call(aux=HCQInfo(name, Estimates(), devs, "COMPUTE:0")) + def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]], - tracker:HCQDepsTracker) -> tuple[list[UOp], set[int]]: + tracker:HCQDepsTracker, slots:dict[str, int]) -> tuple[list[UOp], list[UOp], set[int]]: # collect all buffers which belong to devices dev_bufs:dict[str, dict[int, Any]] = collections.defaultdict(dict) for call, devices in batch: for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)): for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b - zero, n, submits, bumps, waited = UOp.const(0, dtypes.int), len(batch_info), [], [], set() + n, fences, fins, waited = len(batch_info), [], [], set() for _, devgroup in itertools.groupby(sorted(dev_bufs), key=lambda d: d.split(":")[0]): devs = tuple(devgroup) # to finalize the batch, sync all accesses from other devices to buffers that belong to this device fin_deps = [dl for dl in _get_deps(tracker, [list(dev_bufs[d].values()) for d in devs], None, key=(devs, "COMPUTE:0", n)) if dl[0][2] < n] - waits, cur_waited = _build_wait_cmds(fin_deps, devs, "COMPUTE:0") + waits, cur_waited = _build_wait_cmds(slots, fin_deps, devs, "COMPUTE:0") waited |= cur_waited - # wait the syncs, store the device epoch - store = UOp(Ops.INS, arg="store", src=(make_signal(devs), (tl:=make_signal_value(devs)).index(zero) + n)) - submits.append((devs, make_submit(*waits, store, devs=devs, queue="COMPUTE:0"))) - upd = [(tl, n + 1)] + [(make_signal_value(devs, queue=qn), n) - for qn in dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)]) if qn != "COMPUTE:0"] - bumps.append((devs, UOp.barrier(*[s.index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd]))) + # wait the syncs and signal the device epoch, then bump the timeline on the host + timeline, tl = make_signal(devs, tag="timeline_signal"), make_signal(devs, tag="timeline_value") + submit = make_submit(*waits, UOp(Ops.INS, arg="store", src=(timeline, tl.index(0))), devs=devs, queue="COMPUTE:0") + cur = (bump:=tl.after(submit).index(0)).load() + bumps = [bump.store(cur + 1)] - # NOTE: submit before bumps - fins = [UOp.custom_function("hcq", b.sink()).call(aux=HCQInfo("hcq_finalizer", Estimates(), devs, "COMPUTE:0")) for devs, b in submits + bumps] - return fins, waited + # devices running the batch reset their queue signals before each run, fencing on the epoch kept from the previous one + if qs:=dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)]): + prev = make_signal(devs, next(UOp.unique_num)) + fences.append(_hcq_call(devs, "hcq_fence", make_fence(timeline, prev, [make_signal(devs, slots[q]) for q in qs]))) + bumps.append(prev.after(submit).index(0).store(cur)) + fins.append(_hcq_call(devs, "hcq_finalizer", UOp.sink(*bumps))) + return fences, fins, waited def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]: batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch] # schedule deps waited:set[int] = set() + slots:dict[str, int] = collections.defaultdict(lambda: next(UOp.unique_num)) deps_tracker = HCQDepsTracker() call_waits:list[list[UOp]] = [] for tag, ((call, _), (devices, queue)) in enumerate(zip(batch, batch_info)): deps = _get_deps(deps_tracker, _get_call_bufs_by_lane(call, devices), get_call_outs_ins(call)[0], key=(devices, queue, tag)) - cmds, cur_waited = _build_wait_cmds(deps, devices, queue) + cmds, cur_waited = _build_wait_cmds(slots, deps, devices, queue) call_waits.append(cmds) waited |= cur_waited - # build finalizers - finalizers, finalizer_waited = _build_finalizers(batch, batch_info, deps_tracker) + # build fences and finalizers + fences, finalizers, finalizer_waited = _build_finalizers(batch, batch_info, deps_tracker, slots) waited |= finalizer_waited src = [] for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)): - # first queue use, sync prior device work with main signal + # first queue use, sync prior device work with the device timeline if batch_info.index((devices, queue)) == tag: - q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_signal(devices), make_signal_value(devices).index(0) - 1))] + q + epoch = make_signal(devices, tag="timeline_value").index(0) - 1 + q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_signal(devices, tag="timeline_signal"), epoch))] + q # and make hcq call info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), devices, queue) q += [call.replace(arg=replace(call.arg, aux=info))] - # signal queue timeline if someone waits for us - if tag in waited: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, queue), make_signal_value(devices, queue).index(0) + tag))] + # signal the queue if someone waits for us + if tag in waited: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))] src.append(UOp.custom_function("hcq", make_submit(*q, devs=devices, queue=queue).sink()).call(name="hcq", aux=info)) - return src + finalizers + return fences + src + finalizers def sched_hcq_batches(l:UOp) -> UOp: srcs:list[UOp] = [] @@ -216,7 +232,7 @@ def merge_queues(linear:UOp) -> UOp: limits:dict[tuple[tuple[str, ...], str], int] = collections.defaultdict(lambda: JIT_BATCH_SIZE.value) for call in linear.src: - if not isinstance(info:=call.arg.aux, HCQInfo) or info.name == "hcq_finalizer": # non-hcq call or finalizer: close all open queues + if not isinstance(info:=call.arg.aux, HCQInfo) or info.name.startswith("hcq_"): # non-hcq call, fence or finalizer: close all open queues new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call] continue @@ -246,21 +262,11 @@ pm_encode_cmdbufs = PatternMatcher([ # ***************** -def is_value_known_at_link(val:UOp) -> bool: - runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)] - addressed_bufs = [b for g in val.toposort() if g.op is Ops.GETADDR for b in unwrap_mstack(g.buf_uop)] +def get_getaddrs(p:UOp) -> list[UOp]: return [u for u in p.toposort(gate=lambda u: u.op is not Ops.AFTER) if u.op is Ops.GETADDR] - # addr of input params is not known at link time - return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs) - -def is_link_patch(p:UOp, jit:bool) -> bool: - if p.tag == "link": return True - store = p.src[0] if (is_binary_patch:=(p.op is Ops.END and p.src[0].op is Ops.STORE)) else p - if not jit: return store.buf_uop.tag == "program" - return is_binary_patch or (store.op is Ops.STORE and is_value_known_at_link(store.src[1])) - -def trim_link_patches(ctx:tuple[bool, list[UOp]], a:UOp) -> UOp|None: - links, kept = partition(a.src[1:], lambda p: is_link_patch(p, ctx[0])) +def trim_link_patches(ctx:tuple[list[UOp], list[UOp]], a:UOp) -> UOp|None: + links, kept = partition(a.src[1:], lambda p: p.tag == "link") + ctx[0].extend(kept) # keep all patches from the link-time patches' subtrees in the C code afters = [u for u in UOp.sink(*links).toposort() if u.op is Ops.AFTER] @@ -268,18 +274,7 @@ def trim_link_patches(ctx:tuple[bool, list[UOp]], a:UOp) -> UOp|None: return a.src[0].after(*kept, *[d for p in afters for d in p.src[1:]]) if links else None pm_trim_link_patches = PatternMatcher([(UPat(Ops.AFTER, src=(UPat((Ops.PARAM, Ops.MSTACK)),), allow_any_len=True, name="a"), trim_link_patches)]) -def split_patches(ctx:bool, call:UOp) -> UOp|None: - lt_patches:list[UOp] = [] - body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(ctx, lt_patches), name=f"trim link-time patches ({call.arg.aux.name})") - - lt_srcs = collections.defaultdict(list) - for p in lt_patches: lt_srcs[p.buf_uop].append(p) - return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()])) -pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)]) - -# ***************** - -def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]: +def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[UOp, dict[UOp, UOp], tuple[UOp, ...], dict[UOp, int]]: bare = {g: g.replace(src=(g.src[0].without_after,)) for g in gaddrs} order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, repr(b.tag))) @@ -287,26 +282,47 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp table = UOp.placeholder((len(order),), dtypes.uint64, next(UOp.unique_num), device=call.arg.aux.device).rtag(name) reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(slots[bare[g]], dtypes.int)).load() for g in gaddrs} - return reads, (table.after(make_patches(table, [(i * table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else () + fills = (table.after(*make_patches(table, [(i*table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else () + return table, reads, fills, {g:slots[bare[g]] for g in gaddrs} -def make_blob_bufs(call:UOp, blobs:list[UOp]) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]: - bufs = {b: UOp.placeholder((b.max_numel(),), b.dtype, next(UOp.unique_num), device=call.arg.aux.device).rtag("template") for b in blobs} - return bufs, tuple(buf.after(make_binary_patch(buf, b.src[0].arg)) for b,buf in bufs.items()) +def make_scatter_loop(patches:list[UOp], inputs_table:tuple, lt_patches:list[UOp]) -> dict[UOp, UOp]: + (table, _, _, slots), dst, data, subs = inputs_table, patches[0].buf_uop, [], {} + for p in patches: + words = [(off, val, get_getaddrs(val)) for off,val in zip(p.src[0].src[1].src, p.src[1].src)] + data += [off.val << 32 | slots[gaddrs[0]] for off,_,gaddrs in words if gaddrs][::2] + scalars = [(off.val*dst.dtype.itemsize, val) for off,val,gaddrs in words if not gaddrs] + subs[p] = UOp.group(*make_patches(dst, scalars)) if scalars else UOp(Ops.NOOP) -def rm_rt_uops(call:UOp) -> UOp|None: - if not (rt_uops:=[u for u in call.src[0].toposort() if u.op is Ops.GETADDR or (u.op is Ops.BITCAST and u.src[0].op is Ops.BINARY)]): return None - gaddrs, blobs = partition(rt_uops, lambda u: u.op is Ops.GETADDR) - inputs, internals = partition(gaddrs, lambda g: all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop))) + # plan entry: dst word offset << 32 | addr table slot + plan = UOp.placeholder((len(data),), dtypes.uint64, next(UOp.unique_num), device=dst.device).rtag("systems") + entry = plan.index(ridx:=UOp.range(len(data), next(UOp.unique_num), dtype=dtypes.int, src=(plan, dst))).load() + slot, widx = ((entry & 0xffffffff) % table.max_numel()).cast(dtypes.int), ((entry >> 32) % (dst.max_numel()-1)).cast(dtypes.int) # CHECK_OOB bounds + loop = UOp.group(*[dst.index(widx+i).store((table.index(slot).load() >> 32*i).cast(dtypes.uint32)) for i in range(2)]).end(ridx) + lt_patches.append(make_binary_patch(plan, struct.pack(f'<{len(data)}Q', *data))) + subs[patches[0]] = UOp.group(loop, subs[patches[0]]) + return subs + +def is_input_addr(g:UOp) -> bool: return all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop)) + +def split_patches(call:UOp) -> UOp|None: + rt_patches:list[UOp] = [] + lt_patches:list[UOp] = [] + body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(rt_patches, lt_patches), name=f"trim link-time patches ({call.arg.aux.name})") + + # split patches + inputs, internals = partition(dedup(g for p in rt_patches for g in get_getaddrs(p)), is_input_addr) runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop))) + tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))] + reads, fills = {k:v for _,r,_,_ in tables for k,v in r.items()}, [f for t in tables[1:] for f in t[2]] # inputs table is filled by exec + input_patches = [p for p in rt_patches if (gs:=get_getaddrs(p)) and all(map(is_input_addr, gs))] + scatter = make_scatter_loop(input_patches, tables[0], lt_patches) if input_patches else {} + body = body.substitute({p:p.substitute(scatter | reads) for p in rt_patches}) - # exec fills the inputs table with the input addresses every run, so it has no fill patches - (reads, _), *tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))] + \ - [make_blob_bufs(call, blobs)] - reads, fills = reads | {k:v for r,_ in tables for k,v in r.items()}, [f for _,fs in tables for f in fs] - return call.replace(src=(call.src[0].substitute(reads), *call.src[1:], *fills), - arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(g.buf_uop.arg.slot for g in inputs)))))) -pm_rm_rt_uops = PatternMatcher([ - (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_uops)]) + lt_srcs = collections.defaultdict(list) + for p in lt_patches: lt_srcs[p.buf_uop].append(p) + return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills), + arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(b.arg.slot for g in inputs for b in unwrap_mstack(g.buf_uop))))))) +pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)]) # ***************** @@ -372,15 +388,15 @@ def callify_hcq(call:UOp, cf:UOp) -> UOp: pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=( UPat(Ops.CUSTOM_FUNCTION, arg="hcq_args", src=(UPat(Ops.SINK),), name="cf"),), name="call", allow_any_len=True), callify_hcq)]) -hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {} +hcq_compile_cache:dict[bytes, UOp] = {} -@track_rewrites(lambda linear,input_uops,jit,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}") -def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp: +@track_rewrites(lambda linear,input_uops,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}") +def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> UOp: if input_uops is not None: slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))} linear = graph_rewrite(linear, pm_replace_buffers, ctx=(input_uops, slots), walk=True, name="replace buffer") - if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, jit)))) is None: + if (final_linear:=(hcq_compile_cache.get(cache_key:=linear.key))) is None: # prep linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True) linear = graph_rewrite(linear, pm_insert_copy_staging+pm_flatten_linear, name="insert copy staging") @@ -391,11 +407,12 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp: # lowering to hcq ir linear = graph_rewrite(linear, pm_encode_cmdbufs+pm_pack_placeholders, walk=True, name="encode and pack", enter_calls=True) - # patches - linear = graph_rewrite(linear, pm_split_patches+pm_early_simplify+symbolic, ctx=jit, bottom_up=False, name="simplify patches", enter_calls=True) + # patches and runtime uops + linear = graph_rewrite(linear, pm_early_simplify+symbolic, bottom_up=False, name="simplify patches", enter_calls=True) + linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches") # and compile it - linear = graph_rewrite(linear, pm_replace_params, bpm=pm_rm_rt_uops, name="replace rt uops and params") + linear = graph_rewrite(linear, pm_replace_params, name="replace params") final_linear = hcq_compile_cache[cache_key] = graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True) return final_linear @@ -455,11 +472,11 @@ pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: pani def link_buf_key(a:UOp): return a.key, to_tuple(a.device) link_buf_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {} -link_linear_cache:dict[tuple[bytes, bool], UOp] = {} +link_linear_cache:dict[bytes, UOp] = {} -@track_rewrites(lambda _,jit,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}") -def hcq_link(linear:UOp, jit=False, cache=True) -> UOp: - if (linked:=link_linear_cache.get(linear_key:=(linear.key, jit))) is not None: return linked +@track_rewrites(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}") +def hcq_link(linear:UOp, cache=True) -> UOp: + if (linked:=link_linear_cache.get(linear_key:=linear.key)) is not None: return linked bufs = {(j,i):a for j,c in enumerate(linear.src) for i,a in enumerate(c.src[1:], 1) if a.op is Ops.AFTER and unwrap_mstack(a.src[0])[0].tag in HCQ_CACHE_TAGS} @@ -480,7 +497,10 @@ class HCQ2Compiled(Compiled): self.device_id:int = int(device.split(":")[1]) if ":" in device else 0 self.pm_bufferize = PatternMatcher([ - (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].timeline_signal("sentinel", (1 << 64) - 1)), + (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)), + (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].signal("timeline")), + (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].signal("value", 1)), + (UPat(Ops.PARAM, tag="signal", name="b"), lambda ctx, b: ctx[0].signal(b.arg.slot)), (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if b.tag is None else ctx[0].new_buffer(b, cache=ctx[1])) ]) @@ -495,21 +515,15 @@ class HCQ2Compiled(Compiled): return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128)) @functools.cache - def timeline_signal(self, queue:str, init_value:int=0) -> Buffer: + def signal(self, name:str|int, init_value:int=0) -> Buffer: buf = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True) buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value return buf - @functools.cache - def timeline_value(self, queue:str, init_value:int=1) -> Buffer: - buf = Buffer("CPU", 1, dtypes.uint64, preallocate=True) - buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value - return buf - def synchronize(self, timeout:int|None=None): if not hasattr(self, 'iface'): return - sig = self.timeline_signal("COMPUTE:0").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q') - tl = self.timeline_value("COMPUTE:0").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q') + sig = self.signal("timeline").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q') + tl = self.signal("value", 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q') st = time.perf_counter() while sig[0] < tl[0] - 1: if time.perf_counter() - st > (timeout or 3000) / 1000: self.on_device_hang()