diff --git a/extra/hcq2/hcq2.py b/extra/hcq2/hcq2.py index a005b8ad37..bd62b1abfe 100644 --- a/extra/hcq2/hcq2.py +++ b/extra/hcq2/hcq2.py @@ -9,131 +9,29 @@ from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, g from tinygrad.uop.symbolic import symbolic from tinygrad.dtype import dtypes, truncate from tinygrad.runtime.support.hcq import MMIOInterface +from tinygrad.runtime.support.memory import BumpAllocator from tinygrad.renderer import Renderer, Estimates from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear from tinygrad.engine.jit import DepsTracker HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled') -class HCQ2Compiled(Compiled): - timestamp_divider: float = 1000.0 # GPU timestamp counter ticks per microsecond; override per device - - def __init__(self, device:str, allocator:'HCQAllocator', compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None): - self.device_id:int = int(device.split(":")[1]) if ":" in device else 0 - - # default pm bufferize - self.pm_bufferize = PatternMatcher([ - (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx.timeline_signal()), - (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx.timeline_value()), - (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx.timeline_signal("sentinel", (1 << 64) - 1)), - (UPat(Ops.PARAM, name="b"), lambda ctx, b: - Buffer(ctx.device, b.max_numel(), b.dtype, options=BufferSpec(host=False, uncached=True, cpu_access=True, nolru=True)) - if b.tag is not None else None), # TODO: remove nolru - ]) - - super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch) - - @functools.cache - def timeline_signal(self, queue:str|None=None, init_value:int=0) -> Buffer: - buf = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True) - buf._buf.cpu_view().mv.cast('Q')[0] = init_value - return buf - - @functools.cache - def timeline_value(self, queue:str|None=None, init_value:int=1) -> Buffer: - buf = Buffer("CPU", 1, dtypes.uint64, preallocate=True) - buf.as_memoryview(force_zero_copy=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()._buf.cpu_view().mv.cast('Q') - tl = self.timeline_value().as_memoryview(force_zero_copy=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() - - def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent. - - def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1 - - def _select_iface(self): - assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "", \ - f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead" - assert hasattr(self, "ifaces"), "must have ifaces to select an iface" - t = DEV.target(dev:=type(self).__name__[:-6]) - filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}") - filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fall back to mock ifaces - return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in filtered], - f"No interface for {dev}:{self.device_id} is available") - - def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU" - - def finalize(self): - try: self.synchronize() # try to finalize the device in any case - except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}") - - # if the device has an interface, call device_fini to clean up resources - if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini() - -class HCQ2Buffer: - def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQ2Buffer|None=None, view:MMIOInterface|None=None, owner:HCQ2Compiled|None=None): - self.va_addr, self.size, self.meta, self._base, self.view, self.owner = va_addr, size, meta, _base, view, owner - - def offset(self, offset:int=0, size:int|None=None) -> HCQ2Buffer: - return HCQ2Buffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, meta=self.meta, - _base=self._base or self, view=(self.view.view(offset=offset, size=size) if self.view is not None else None)) - - def cpu_view(self) -> MMIOInterface: - assert self.view is not None, "buffer has no cpu_view" - return self.view - - @property - def base(self) -> HCQ2Buffer: return self._base or self - -class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]): - def _map(self, buf:HCQ2Buffer) -> HCQ2Buffer: - if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented") - return self._do_map(buf) - - @suppress_finalizing - def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None): - self.dev.synchronize() - if options is not None and options.external_ptr is not None: return - if hasattr(self, '_do_free'): self._do_free(buf, options) - - def _unmap(self, mb): - self.dev.synchronize() - self.dev.iface.free(mb) - - def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size) - - def _wrap(self, dev:str, sz:int, opaque:HCQ2Buffer) -> Buffer: - return Buffer(dev, sz, dtypes.uint8, opaque=opaque, options=BufferSpec(external_ptr=1)) - - def _copy(self, dst:Buffer, src:Buffer): - from tinygrad.engine.realize import run_linear - su = UOp.from_buffer(src) - run_linear(UOp(Ops.LINEAR, dtypes.void, (su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), update_stats=False) - - def _copyin(self, dest:HCQ2Buffer, src:memoryview): - s = Buffer(self.dev.device, len(src), dtypes.uint8, options=BufferSpec(host=True), preallocate=True) - s._buf.cpu_view()[:len(src)] = src - self._copy(self._wrap(self.dev.device, len(src), dest), s) - - def _copyout(self, dest:memoryview, src:HCQ2Buffer): - d = Buffer(self.dev.device, len(dest), dtypes.uint8, options=BufferSpec(host=True), preallocate=True) - self._copy(d, self._wrap(self.dev.device, len(dest), src)) - self.dev.synchronize() - dest[:] = d._buf.cpu_view()[:len(dest)] - - # def _as_buffer(self, buf): return buf.cpu_view().mv - # ***************** # 0. helpers HCQ_DEVS = frozenset(("AMD",)) HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",)) +HCQ_CACHE_TAGS = frozenset(("program", "systems", "template")) + +@dataclass(frozen=True) +class HCQInfo: + name:str + estimates:Estimates + device:tuple[str, ...] + queue:str + + input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call + inputs:int|None = None def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c @@ -152,28 +50,23 @@ 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) 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 "buf") + return UOp.param(next(UOp.unique_num) if unique else 0, dtype, shape=(size,), device=devs).rtag(name or "temp") -def tag_patch(val:UOp) -> str: - if any(u.op in (Ops.LOAD, Ops.INDEX) for u in val.toposort()): return "rt" - if not (gaddrs:=[u for u in val.toposort() if u.op is Ops.GETADDR]): return "any" - return "rt" if any(x.op is Ops.PARAM and x.tag is None for g in gaddrs for x in unwrap_mstack(g.buf_uop)) else "link" +def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp: + return buf.index(UOp.const(dtypes.int, off // buf.dtype.itemsize)).store(val.simplify().cast(dtype or buf.dtype)) -def make_patch(buf:UOp, off:sint, val:UOp, dtype=None, tag=None) -> UOp: - return buf.index(UOp.const(dtypes.int, off // buf.dtype.itemsize)).store(val.simplify().cast(dtype or buf.dtype)).rtag(tag or tag_patch(val)) - -def make_binary_patch(buf:UOp, blob:bytes, tag=None) -> UOp: - data, dt, isz = UOp(Ops.BINARY, dtypes.uint8, src=(), arg=blob), buf.dtype, buf.dtype.itemsize +def make_binary_patch(buf:UOp, blob:bytes) -> UOp: + data, isz = UOp(Ops.BINARY, dtypes.uint8, src=(), arg=blob), buf.dtype.itemsize r = UOp.range(len(blob) // isz, next(UOp.unique_num)) - return buf.index(r).store(UOp(Ops.BITCAST, buf.dtype, (data,)).index(r).load()).end(r).rtag(tag or "any") + return buf.index(r).store(UOp(Ops.BITCAST, buf.dtype, (data,)).index(r).load()).end(r) def make_cmdbuf(lin, devs): blob, patches = b'', [] for s in (s for ins in lin.src for s in ins.src): if (ssimp:=s.simplify()).op is not Ops.CONST: patches.append((len(blob), ssimp)) blob += struct.pack(f'<{ssimp.dtype.fmt}', ssimp.arg if ssimp.op is Ops.CONST else 0x0) - buf = make_placeholder(devs, len(blob) // 4, dtypes.uint32) - return buf.after(make_binary_patch(buf, blob, tag="link"), *[make_patch(buf, off, s) for off, s in patches]) + cmdbuf = make_placeholder(devs, len(blob) // 4, dtypes.uint32, name="cmdbuf") + return cmdbuf.after(make_binary_patch(cmdbuf, blob), *[make_patch(cmdbuf, off, s) for off, s in patches]) def make_mstack(uops): return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, tuple(uops)) @@ -183,24 +76,14 @@ def make_signal_value(devs, queue=None): return make_placeholder(devs, 1, dtypes.uint64, (queue, "timeline_value") if queue else "timeline_value", unique=False) def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp: - return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, dtypes.void, src=tuple(cmds), arg=(to_tuple(devs), queue))) + return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, src=tuple(cmds), arg=(to_tuple(devs), queue))) def get_submit(ast:UOp) -> UOp: return next(u for u in ast.toposort() if u.op is Ops.CUSTOM_FUNCTION and u.arg == "submit_cmdbuf") -@dataclass(frozen=True) -class HCQInfo: - name:str - estimates:Estimates - device:tuple[str, ...] - queue:str - - input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call - inputs:int|None = None - # ***************** # 0.1. prep: replace buffers with params def replace_call_buffers(ctx:list[UOp], call:UOp) -> UOp|None: - ctx += [s for s in dedup(call.src[1:]) if s not in ctx and s.op not in (Ops.PARAM, Ops.BIND)] + ctx += [s for s in call.src[1:] if s not in ctx and s.op not in (Ops.PARAM, Ops.BIND)] return call.replace(src=call.src[:1] + tuple(s if s.op in (Ops.PARAM, Ops.BIND) else s.param_like(ctx.index(s)) for s in call.src[1:])) pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_buffers)]) @@ -213,7 +96,7 @@ def stage_copy(dst:UOp, src:UOp) -> UOp|None: if not (_need_staging(src, dst) or _need_staging(dst, src)): return None stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.itemsize, dtypes.uint8) - return UOp(Ops.LINEAR, dtypes.void, (src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage))) + return UOp(Ops.LINEAR, src=(src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage))) pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)]) # ***************** @@ -225,8 +108,7 @@ def tag_hcq_call(ctx:itertools.count, call:UOp) -> UOp: queue = "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0" info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), to_tuple(hcq_devs), queue) return call.replace(arg=replace(call.arg, aux=info)).rtag(next(ctx)) -pm_tag_hcq_calls = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), - lambda ctx, linear: linear.replace(src=tuple(tag_hcq_call(ctx, s) for s in linear.src)))]) +pm_tag_hcq_calls = PatternMatcher([(UPat(Ops.LINEAR, name="l"), lambda ctx, l: l.replace(src=tuple(tag_hcq_call(ctx, s) for s in l.src)))]) # ***************** # 2.2. deps tracking @@ -378,7 +260,7 @@ def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp: data, info = prg.arg buf = make_placeholder(devs, data.kernargs_alloc_size // 4, dtypes.uint32, name="kernargs") words = [w for gi in info.globals for w in data64_le(make_getaddr(get_call_arg_uops(call)[gi], devs))] + list(info.vars) - return buf.after(*[make_patch(buf, i * 4, w, tag="rt") for i, w in enumerate(words)]) + return buf.after(*[make_patch(buf, i * 4, w) for i, w in enumerate(words)]) # ***************** # 4.2. hcq lowering: ops to ir @@ -390,7 +272,18 @@ pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbu # ***************** -def is_link_patch(p:UOp, jit:bool) -> bool: return p.tag == "link" or p.tag == "any" # TODO: jit +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 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: + store = p.src[0] if (is_binary_patch:=p.op is Ops.END) 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])) @@ -410,26 +303,38 @@ pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION # ***************** -def _make_getaddrs_sub(call:UOp, gaddrs:list[UOp], name:str): +def make_addr_table(call:UOp, gaddrs:list[UOp], name:str): bare = {g: g.replace(src=(unwrap_after(g.src[0]),)) for g in gaddrs} order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, to_tuple(b.tag))) - b = make_placeholder(call.arg.aux.device, len(order), dtypes.uint64, name) + slots, table = {g:i for i,g in enumerate(order)}, make_placeholder(call.arg.aux.device, len(order), dtypes.uint64, name) - sub = {g: b.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, order.index(gr))).load() for g,gr in bare.items()} - return sub, (b.after(*[make_patch(b, i * b.dtype.itemsize, gr) for i,gr in enumerate(order)]),) if order else () + reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs} + return reads, (table.after(*[make_patch(table, i * table.dtype.itemsize, addr) for addr, i in slots.items()]),) if slots else () def rm_rt_getaddrs(call:UOp) -> UOp|None: if not (gaddrs:=[u for u in call.src[0].toposort() if u.op is Ops.GETADDR]): return None - inputs, systems = partition(gaddrs, lambda g: all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop))) + 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))) + runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop))) - (inpsub, _), (syssub, sysarg) = _make_getaddrs_sub(call, inputs, "inputs"), _make_getaddrs_sub(call, systems, "systems") - return call.replace(src=(call.src[0].substitute(inpsub | syssub), *call.src[1:], *sysarg), + # exec fills the inputs table with the input addresses every run, so it has no fill patches + (input_reads, _), (rt_reads, rt_fills), (sys_reads, sys_fills) = (make_addr_table(call, gs, name) for gs, name in + ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))) + return call.replace(src=(call.src[0].substitute(input_reads | rt_reads | sys_reads), *call.src[1:], *rt_fills, *sys_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_getaddrs = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_getaddrs)]) # ***************** +def rm_rt_binaries(call:UOp) -> UOp|None: + if not (blobs:=[u for u in call.src[0].toposort() if u.op is Ops.BITCAST and u.src[0].op is Ops.BINARY]): return None + blob_bufs = {blob: make_placeholder(call.arg.aux.device, blob.max_numel(), blob.dtype, "template") for blob in blobs} + fills = [buf.after(make_binary_patch(buf, blob.src[0].arg)) for blob, buf in blob_bufs.items()] + return call.replace(src=(call.src[0].substitute(blob_bufs), *call.src[1:], *fills)) +pm_rm_rt_binaries = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_binaries)]) + +# ***************** + def replace_params(call:UOp) -> UOp|None: body, variables, param_ops = call.src[0], call.src[0].variables(), {Ops.PARAM, Ops.MSTACK} args = dedup([s for u in body.toposort(gate=lambda u: u.op not in param_ops) for s in u.src if s.op in param_ops and s not in variables]) @@ -447,35 +352,32 @@ pm_replace_params = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTIO # ***************** def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp: + base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()) itemsize = bv.src[0].dtype.itemsize if unwrap_after(bv.src[0]).op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize - return UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0],), arg=g.arg) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize) + return UOp(Ops.GETADDR, dtypes.uint64, src=(base,), arg=g.arg) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize) pm_early_simplify = PatternMatcher([ - # getaddr(slice(base, off)) -> getaddr(base) + byte offset - (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"),), name="g"), resolve_getaddr_slice), + (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice), + (UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"),), allow_any_len=True, name="x"), + lambda bv,x: x.replace(src=(bv.src[0], x.src[1] + bv.src[1].cast(x.src[1].dtype), *x.src[2:]))), ]) # ***************** # 5.3. pack placeholders buffers -# def pack_hcq_placeholders(call:UOp) -> UOp|None: -# bufs = [b for b in call.src[0].toposort() if b.op is Ops.PARAM and b.tag in (maxtags:={"scratch"}) | (sumtags:={"program", "kernargs"})] - -# off_per_buf:dict[UOp, int] = {} -# size_per_tag:dict[str, int] = {} -# for b in bufs: -# bsz = b.max_numel() -# if b.tag in maxtags: size_per_tag[b.tag] = max(size_per_tag.get(b.tag, 0), bsz) -# elif b.tag in sumtags: -# off_per_buf[b] = round_up(size_per_tag.get(b.tag, 0), {"program": 0x1000}.get(b.tag, 128)) -# size_per_tag[b.tag] = off_per_buf[b] + bsz - -# count_per_tag = collections.Counter(b.tag for b in bufs) -# ref_bufs = {b.tag:b for b in bufs if count_per_tag[b.tag] > 1} -# bases = {tag:UOp.new_buffer(b.device, size_per_tag[tag], b.dtype).rtag(tag) for tag,b in ref_bufs.items()} -# subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.index, off_per_buf.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases} -# return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None -# pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)]) +def pack_hcq_placeholders(call:UOp) -> UOp|None: + bufs = [b for b in call.src[0].toposort() if b.op is Ops.PARAM and b.tag in {"scratch", "kernargs"}] + offs, sizes = {}, {} + for b in bufs: + if b.tag == "scratch": sizes[b.tag] = max(sizes.get(b.tag, 0), b.max_numel()) + else: + offs[b] = round_up(sizes.get(b.tag, 0), 128 // b.dtype.itemsize) + sizes[b.tag] = offs[b] + b.max_numel() + counts = collections.Counter(b.tag for b in bufs) + bases = {b.tag:make_placeholder(b.device, sizes[b.tag], b.dtype, b.tag) for b in bufs if counts[b.tag] > 1} + subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.index, offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases} + return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None +pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)]) # ***************** # 8. callify hcq programs @@ -483,13 +385,13 @@ pm_early_simplify = PatternMatcher([ pm_callify_hcq = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="hcq", src=(UPat(Ops.SINK),), name="cf"), lambda cf: cf.replace(src=(to_program(cf.src[0].replace(arg=KernelInfo("hcq_submit"), tag=1), Device["CPU"].renderer),)))]) -hcq_compile_cache:dict[bytes, UOp] = {} +hcq_compile_cache:dict[tuple[bytes, bool], 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: +@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: if input_uops is not None: linear = graph_rewrite(linear, pm_replace_buffers, ctx=input_uops, walk=True, enter_calls=True, name="replace buffer") - if (final_linear:=(hcq_compile_cache.get(cache_key:=linear.key))) is None: + if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, jit)))) is None: # schedule 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") @@ -504,14 +406,15 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> UOp: linear = graph_rewrite(linear, pm_add_inner_loads, ctx=(waited:=set()), walk=True, name="add loads", enter_calls=True) linear = graph_rewrite(linear, pm_add_inner_stores, ctx=waited, walk=True, name="add stores", enter_calls=True) linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True) + linear = graph_rewrite(linear, pm_pack_placeholders, walk=True, name="pack placeholders") # pie - linear = graph_rewrite(linear, pm_split_patches, ctx=input_uops is None, walk=True, name="split rt/lt patches") + linear = graph_rewrite(linear, pm_split_patches, ctx=jit, walk=True, name="split rt/lt patches") + linear = graph_rewrite(linear, pm_early_simplify + symbolic, bottom_up=False, name="simplify packed placeholders", enter_calls=True) linear = graph_rewrite(linear, pm_rm_rt_getaddrs, walk=True, name="replace rt getaddrs") + linear = graph_rewrite(linear, pm_rm_rt_binaries, walk=True, name="replace rt binaries") linear = graph_rewrite(linear, pm_replace_params, walk=True, name="replace with args") - linear = graph_rewrite(linear, pm_early_simplify + symbolic, bottom_up=False, name="early simplify patches", enter_calls=True) - # and compile it final_linear = hcq_compile_cache[cache_key] = graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True) @@ -520,9 +423,9 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> UOp: # ***************** # 6. bufferize placeholders: replace placeholders with real buffers. -def bufferize_buf(buf:UOp) -> UOp|None: +def bufferize_buf(ctx:bool, buf:UOp) -> UOp|None: if buf.tag is None: return None - return make_mstack(tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=dv), "CPU") for dev in to_tuple(buf.device))) + return make_mstack(tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=(dv, ctx)), "CPU") for dev in to_tuple(buf.device))) pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)]) # ***************** @@ -567,8 +470,140 @@ pm_resolve_patches = PatternMatcher([ pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: panic(RuntimeError, f"AFTER left at hcq_link: {a.src[0].op}"))]) -@track_rewrites(lambda _,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}") -def hcq_link(linear:UOp) -> UOp: - linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, walk=True, name="bufferize placeholders") +hcq_link_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {} + +def link_cache_key(a:UOp): return a.key, to_tuple(a.device) +pm_link_cache = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: hcq_link_cache.get(link_cache_key(a)))]) + +@track_rewrites(lambda _,jit,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}") +def hcq_link(linear:UOp, jit=False) -> UOp: + cacheable = {(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} + hits = {a.src[0]:hcq_link_cache[key] for a in cacheable.values() if (key:=link_cache_key(a)) in hcq_link_cache} + linear = graph_rewrite(linear, pm_link_cache, name="apply link cache").substitute(hits, walk=True) + linear = graph_rewrite(linear, pm_bufferize, ctx=jit, bottom_up=True, walk=True, name="bufferize placeholders") linear = graph_rewrite(linear, pm_resolve_patches + symbolic, bottom_up=False, name="simplify patches") - return graph_rewrite(linear, pm_assert_no_afters, name="assert no afters") + linear = graph_rewrite(linear, pm_assert_no_afters, name="assert no afters") + for (j,i),a in cacheable.items(): hcq_link_cache.setdefault(link_cache_key(a), linear.src[j].src[i]) + return linear + +# ***************** +# Device classes + +class HCQ2Compiled(Compiled): + timestamp_divider: float = 1000.0 + + def __init__(self, device:str, allocator:HCQAllocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None): + self.device_id:int = int(device.split(":")[1]) if ":" in device else 0 + + self.pm_bufferize = PatternMatcher([ + (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].timeline_signal()), + (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].timeline_value()), + (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].timeline_signal("sentinel", (1 << 64) - 1)), + (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if b.tag is None else ctx[0].new_buffer(b, jit=ctx[1])) + ]) + + super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch) + + self.rt_buffer = Buffer(self.device, 64 << 20, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True)) + self.rt_allocator = BumpAllocator(64 << 20, wrap=False) + + def new_buffer(self, b:UOp, jit:bool) -> Buffer: + if jit or b.tag in HCQ_CACHE_TAGS: return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(cpu_access=True, nolru=True)) + 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|None=None, init_value:int=0) -> Buffer: + buf = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True) + buf._buf.cpu_view().mv.cast('Q')[0] = init_value + return buf + + @functools.cache + def timeline_value(self, queue:str|None=None, init_value:int=1) -> Buffer: + buf = Buffer("CPU", 1, dtypes.uint64, preallocate=True) + buf.as_memoryview(force_zero_copy=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()._buf.cpu_view().mv.cast('Q') + tl = self.timeline_value().as_memoryview(force_zero_copy=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() + + def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent. + + def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1 + + def _select_iface(self): + assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "", \ + f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead" + assert hasattr(self, "ifaces"), "must have ifaces to select an iface" + t = DEV.target(dev:=type(self).__name__[:-6]) + filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}") + filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fall back to mock ifaces + return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in filtered], + f"No interface for {dev}:{self.device_id} is available") + + def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU" + + def finalize(self): + try: self.synchronize() # try to finalize the device in any case + except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}") + + # if the device has an interface, call device_fini to clean up resources + if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini() + +class HCQ2Buffer: + def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQ2Buffer|None=None, view:MMIOInterface|None=None, owner:HCQ2Compiled|None=None): + self.va_addr, self.size, self.meta, self._base, self.view, self.owner = va_addr, size, meta, _base, view, owner + + def offset(self, offset:int=0, size:int|None=None) -> HCQ2Buffer: + return HCQ2Buffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, meta=self.meta, + _base=self._base or self, view=(self.view.view(offset=offset, size=size) if self.view is not None else None)) + + def cpu_view(self) -> MMIOInterface: + assert self.view is not None, "buffer has no cpu_view" + return self.view + + @property + def base(self) -> HCQ2Buffer: return self._base or self + +class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]): + def _map(self, buf:HCQ2Buffer) -> HCQ2Buffer: + if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented") + return self._do_map(buf) + + @suppress_finalizing + def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None): + self.dev.synchronize() + if options is not None and options.external_ptr is not None: return + if hasattr(self, '_do_free'): self._do_free(buf, options) + + def _unmap(self, mb): + self.dev.synchronize() + self.dev.iface.free(mb) + + def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size) + + def _wrap(self, dev:str, sz:int, opaque:HCQ2Buffer) -> Buffer: + return Buffer(dev, sz, dtypes.uint8, opaque=opaque, options=BufferSpec(external_ptr=1)) + + def _copy(self, dst:Buffer, src:Buffer): + from tinygrad.engine.realize import run_linear + su = UOp.from_buffer(src) + run_linear(UOp(Ops.LINEAR, src=(su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), update_stats=False) + + def _copyin(self, dest:HCQ2Buffer, src:memoryview): + s = Buffer(self.dev.device, len(src), dtypes.uint8, options=BufferSpec(host=True), preallocate=True) + s._buf.cpu_view()[:len(src)] = src + self._copy(self._wrap(self.dev.device, len(src), dest), s) + + def _copyout(self, dest:memoryview, src:HCQ2Buffer): + d = Buffer(self.dev.device, len(dest), dtypes.uint8, options=BufferSpec(host=True), preallocate=True) + self._copy(d, self._wrap(self.dev.device, len(dest), src)) + self.dev.synchronize() + dest[:] = d._buf.cpu_view()[:len(dest)] + + # def _as_buffer(self, buf): return buf.cpu_view().mv diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index d5c4ed2a6f..b9d9648dfd 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -280,7 +280,7 @@ def amd_build_program(prg:UOp) -> UOp: kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp, enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER) buf = make_placeholder(prg.device, len(image), dtypes.uint8, "program") - cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image), tag="link")),), arg=(data, prg.arg)) + cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg)) return cached class AMDAllocator(HCQAllocator['AMDDevice']): @@ -580,7 +580,7 @@ class AMDDevice(HCQ2Compiled): # Scratch setup self.max_private_segment_size = 0 - self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.max_numel()))]) + self.pm_bufferize + self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx[0].scratch_buffer(b.max_numel()))]) + self.pm_bufferize self.pmc_enabled:bool = PROFILE > 0 and PMC > 0 if self.pmc_enabled: @@ -630,8 +630,8 @@ class AMDDevice(HCQ2Compiled): self.pm_bufferize = PatternMatcher([ (UPat(Ops.PARAM, tag={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"] ] + [ - (UPat(Ops.PARAM, tag={(qname, "timeline_signal")}), lambda ctx, q=qname: ctx.timeline_signal(q)), - (UPat(Ops.PARAM, tag={(qname, "timeline_value")}), lambda ctx, q=qname: ctx.timeline_value(q)), + (UPat(Ops.PARAM, tag={(qname, "timeline_signal")}), lambda ctx, q=qname: ctx[0].timeline_signal(q)), + (UPat(Ops.PARAM, tag={(qname, "timeline_value")}), lambda ctx, q=qname: ctx[0].timeline_value(q)), ]) + self.pm_bufferize return queue diff --git a/test/backend/test_tensor_variable.py b/test/backend/test_tensor_variable.py index 8943200cc8..37d29266e2 100644 --- a/test/backend/test_tensor_variable.py +++ b/test/backend/test_tensor_variable.py @@ -28,16 +28,16 @@ class TestTensorVariable(unittest.TestCase): def test_variable_tensor_dtype_arg(self): vv = Variable("a", 1, 10).bind(2) - # TODO: dtype arg is silently dropped for a symbolic int, should be honored (or rejected) - try: - self.assertEqual(Tensor(vv, dtype=dtypes.float32).dtype, dtypes.float32) - except AssertionError: pass + t = Tensor(vv, dtype=dtypes.float32) + self.assertEqual(t.dtype, dtypes.float32) + self.assertEqual(t.item(), 2.0) def test_unbound_variable_tensor(self): - # TODO: Tensor creation from unbound variable should assert - # with self.assertRaises(AssertionError): Tensor(Variable("u", 1, 10)) - t = Tensor(Variable("u", 1, 10)) - self.assertRaises(KeyError, t.item) # today it builds silently and fails at execution + # an unbound variable schedules fine, but can't execute + with self.assertRaisesRegex(RuntimeError, "unbound"): Tensor(Variable("u", 1, 10)).item() + with self.assertRaisesRegex(RuntimeError, "unbound"): (Tensor(Variable("u", 1, 10)) + 1).item() + # bound variables in an expression are fine + self.assertEqual(Tensor(Variable("u", 1, 10).bind(2) + 1).item(), 3) def test_shrink_beyond_buffer_variable(self): # TODO: shrink by a variable whose vmax exceeds the dim should fail at build, today only CHECK_OOB=1 rejects it @@ -49,11 +49,9 @@ class TestTensorVariable(unittest.TestCase): # NOTE: the buffer dim must cover the variable's vmax vv = Variable("a", 1, 10).bind(2) self.assertEqual((Tensor.ones(10).contiguous()[:vv] * Tensor(vv)).sum().item(), 4.0) + # a vmin=0 symbolic dim broadcasts too v0 = Variable("z", 0, 10).bind(2) - # TODO: broadcasting a vmin=0 symbolic dim fails, max(dim, 1) cannot be proven equal to dim - try: - self.assertEqual((Tensor.ones(10).contiguous()[:v0] * Tensor(v0)).sum().item(), 4.0) - except IndexError: pass + self.assertEqual((Tensor.ones(10).contiguous()[:v0] * Tensor(v0)).sum().item(), 4.0) def test_inner_tvar_node(self): vv = Variable("w", 0, 10).bind(2) diff --git a/test/null/test_symbolic_tensor.py b/test/null/test_symbolic_tensor.py index 7ac75e71c9..75b6156c92 100644 --- a/test/null/test_symbolic_tensor.py +++ b/test/null/test_symbolic_tensor.py @@ -1,6 +1,23 @@ import unittest from tinygrad import Variable from tinygrad.tensor import Tensor +from tinygrad.uop.ops import _broadcast_shape + +class TestBroadcastShape(unittest.TestCase): + def test_symbolic(self): + v = Variable("v", 1, 10) + self.assertEqual(_broadcast_shape((v,), (1,)), (v,)) + self.assertEqual(_broadcast_shape((v,), ()), (v,)) + self.assertEqual(_broadcast_shape((v,), (v,)), (v,)) + with self.assertRaises(IndexError): _broadcast_shape((v,), (5,)) + + def test_symbolic_vmin_zero(self): + # a symbolic dim that may be 0 still broadcasts against 1 to itself + v0 = Variable("v0", 0, 10) + self.assertEqual(_broadcast_shape((v0,), (1,)), (v0,)) + self.assertEqual(_broadcast_shape((v0,), ()), (v0,)) + self.assertEqual(_broadcast_shape((3, v0), (3, 1)), (3, v0)) + with self.assertRaises(IndexError): _broadcast_shape((v0,), (5,)) class TestSymbolic(unittest.TestCase): def assert_tuple_equal(self, x, y): diff --git a/test/unit/test_multitensor.py b/test/unit/test_multitensor.py index 922dabe15f..b425931132 100644 --- a/test/unit/test_multitensor.py +++ b/test/unit/test_multitensor.py @@ -78,8 +78,9 @@ class TestMultiTensor(unittest.TestCase): self.assertEqual(Y.device, devices_2) np.testing.assert_equal(X.numpy(), Y.numpy()) - with self.assertRaises(AssertionError): - _ = Tensor(X.uop, dtype=dtypes.float) + Z = Tensor(X.uop, dtype=dtypes.float) + self.assertEqual(Z.dtype, dtypes.float) + np.testing.assert_equal(Z.numpy(), [1.0, 2.0]) def test_sharded_arange(self): sharded_arange = Tensor.arange(1000).clone().shard(devices_2, 0) diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index f54060b987..016426492e 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)) + linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value), jit=True) 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 @@ -197,7 +197,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) + def linear(self) -> UOp: return link_linear(self._linear, jit=True) 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 d8397b0709..0a096ed2d9 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -214,7 +214,7 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: buf = b.bufs[j] if isinstance(b:=call.src[1+call.arg.aux.inputs].buffer, MultiBuffer) else b buf.ensure_allocated()._buf.cpu_view().view(fmt='Q')[:len(addrs)] = array.array('Q', addrs) - pm_exec.rewrite(call.replace(src=(ast,) + call.src[1:]), replace(ctx, update_stats=False, wait=True)) + pm_exec.rewrite(call.replace(src=(ast,) + call.src[1:]), replace(ctx, update_stats=False)) for d in call.arg.aux.device: with track_stats(ctx, call, d, [], ctx.var_vals): @@ -259,24 +259,24 @@ pm_exec = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate), ]) -def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None) -> UOp: +def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, jit=False) -> 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"): from extra.hcq2.hcq2 import hcq_compile - linear = hcq_compile(linear, input_uops) + linear = hcq_compile(linear, input_uops, jit=jit) return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True) -def link_linear(linear:UOp) -> UOp: +def link_linear(linear:UOp, jit=False) -> UOp: if getenv("HCQ2"): from extra.hcq2.hcq2 import hcq_link - linear = hcq_link(linear) + linear = hcq_link(linear, jit=jit) return 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)) + if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs, jit=False)) ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2) for call in linear.src: pm_exec.rewrite(call, ctx) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 958ba23dd8..6d8865e1bc 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -69,9 +69,9 @@ class Tensor(RandMixin): # create a UOp from the different types of inputs if isinstance(data, UOp): - assert _dtype is None or _dtype==data.dtype or data.dtype==dtypes.index, f"dtype mismatch: {_dtype} vs {data.dtype}" # if data is dtype.index that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of if data.dtype == dtypes.index: data = _index_to_concrete_int(data) + if _dtype is not None: data = data.cast(_dtype) elif data is None: data = UOp.const(_dtype or dtypes.default_float, 0) elif isinstance(data, get_args(ConstType)): diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 2ad3a0d0f7..65c7f4db75 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -70,11 +70,13 @@ def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]: max_dim = max(len(s) for s in shapes) return tuple((1,)*(max_dim-len(s))+s for s in shapes) def _broadcast_shape(*shapes:tuple[sint, ...]) -> tuple[sint, ...]: - shaped_aligned_left = _align_left(*shapes) - ret = tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*shaped_aligned_left)) - if not all(resolve(s == ns) or resolve(s == 1) for shape in shaped_aligned_left for s,ns in zip(shape, ret)): - raise IndexError(f"shape mismatch: objects cannot be broadcast to a single shape {shapes}") - return ret + # per right-aligned dim: sizes of 1 broadcast to the others, which must all agree + ret = [] + for sizes in zip(*_align_left(*shapes)): + if len(rest:=dedup([s for s in sizes if isinstance(s, UOp) or s != 1])) > 1: + raise IndexError(f"shape mismatch: objects cannot be broadcast to a single shape {shapes}") + ret.append(rest[0] if rest else 1) + return tuple(ret) def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop @@ -1164,7 +1166,9 @@ class ProgramInfo: local_size = tuple([sym_infer(sz, var_vals) for sz in self.local_size]) if self.local_size is not None else None return global_size, local_size - def vals(self, var_vals:dict[str, int]): return tuple(var_vals[k.expr] if k.expr not in self.runtimevars else None for k in self.vars) + def vals(self, var_vals:dict[str, int]) -> tuple[int|None, ...]: + try: return tuple(var_vals[k.expr] if k.expr not in self.runtimevars else None for k in self.vars) + except KeyError as e: raise RuntimeError(f"unbound Variable {e} used by {self.function_name}") from None @staticmethod def from_sink(sink:UOp, aux:tuple=()) -> ProgramInfo: