From 131ded7ecff2831676ca80d582fc37ab34943467 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 24 Jun 2026 18:32:28 +0000 Subject: [PATCH] simpler --- test/null/test_memory_coalescing.py | 51 ----------- tinygrad/codegen/__init__.py | 1 - tinygrad/codegen/late/coalese.py | 128 +++++++++++++--------------- 3 files changed, 57 insertions(+), 123 deletions(-) delete mode 100644 test/null/test_memory_coalescing.py diff --git a/test/null/test_memory_coalescing.py b/test/null/test_memory_coalescing.py deleted file mode 100644 index afd6ea8cb8..0000000000 --- a/test/null/test_memory_coalescing.py +++ /dev/null @@ -1,51 +0,0 @@ -import unittest - -from tinygrad.dtype import dtypes -from tinygrad.helpers import Context, Target -from tinygrad.uop.ops import UOp, KernelInfo -from tinygrad.renderer.cstyle import OpenCLRenderer -from tinygrad.codegen import to_program - - -def render_stores(offsets:list[int], dtype=dtypes.float) -> str: - src = UOp.param(0, dtype.ptr(256)) - out = UOp.param(1, dtype.ptr(256)) - stores = [] - for off in offsets: - idx = UOp.const(dtypes.int, off) - val = src.index(idx, ptr=True).load() - stores.append(out.index(idx, ptr=True).store(val + UOp.const(dtype, 1))) - with Context(IMAGE=1): - arch = "IMAGE_PITCH_ALIGNMENT=64" + (",cl_khr_fp16" if dtype == dtypes.half else "") - prg = to_program(UOp.group(*stores).sink(arg=KernelInfo(opts_to_apply=())), OpenCLRenderer(Target("CL", arch=arch))) - return prg.src[3].arg - - -class TestMemoryCoalescing(unittest.TestCase): - def test_aligned_vec4_uses_images(self): - src = render_stores([0, 1, 2, 3]) - self.assertIn("image2d_t", src) - self.assertIn("read_imagef", src) - self.assertIn("write_imagef", src) - - def test_aligned_half_vec4_uses_images(self): - src = render_stores([0, 1, 2, 3], dtypes.half) - self.assertIn("image2d_t", src) - self.assertIn("read_imagef", src) - self.assertIn("write_imagef", src) - - def test_partial_group_does_not_use_images(self): - src = render_stores([0, 1, 2, 3, 4]) - self.assertNotIn("image2d_t", src) - self.assertNotIn("read_imagef", src) - self.assertNotIn("write_imagef", src) - - def test_misaligned_vec4_does_not_use_images(self): - src = render_stores([1, 2, 3, 4]) - self.assertNotIn("image2d_t", src) - self.assertNotIn("read_imagef", src) - self.assertNotIn("write_imagef", src) - - -if __name__ == "__main__": - unittest.main() diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index e179992feb..904ad7cc70 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -122,7 +122,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # do memory coalesing (late) sink = memory_coalesing(sink, ren) - sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing+gep_pushing, name="lower post-coalescing index dtypes") # instruction selection decompositions pm_decomp = pm_decomp+\ diff --git a/tinygrad/codegen/late/coalese.py b/tinygrad/codegen/late/coalese.py index 5464ff1f46..de5558d1d5 100644 --- a/tinygrad/codegen/late/coalese.py +++ b/tinygrad/codegen/late/coalese.py @@ -6,8 +6,6 @@ from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp from tinygrad.helpers import getenv, IMAGE from tinygrad.renderer import Renderer -_Memory = dict[tuple[Ops, UOp, Any, Any], dict[int, list[UOp]]] - pm_imageh_store = PatternMatcher([ # store(idx, x) is actually store(idx, x.cast(half)) so we can pull the cast into the store (UPat.var("x", dtypes.float).cast(dtypes.half), lambda x: x), @@ -17,64 +15,6 @@ pm_imageh_store = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda x: x.cast(dtypes.float)) ]) -def _grouped_offsets(offsets:dict[int, list[UOp]]) -> list[list[int]]: - return [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])] - -def _valid_idx(idx:UOp) -> tuple[UOp, UOp|None]: - return (idx.src[1], idx.src[0]) if idx.op is Ops.WHERE and idx.src[2].arg is Invalid else (idx, None) - -def _base_offset(idx:UOp) -> tuple[Any, int]: - if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return idx.src[0], idx.src[1].arg - if idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: return idx.src[1], idx.src[0].arg - if idx.op is Ops.CONST and idx.arg is Invalid: return "INVALID", 0 - if idx.op is Ops.CONST: return "CONST", idx.arg - return idx, 0 - -def _offset(base:UOp|str, off:int) -> UOp: return (base+off) if isinstance(base, UOp) else UOp.const(dtypes.int, off) - -def _image_buf(buf:UOp, ctx:Renderer|None) -> UOp|None: - if isinstance(buf.dtype, ImageDType): return buf - if ctx is None or not IMAGE or ctx.target.device not in {"QCOM", "CL", "PYTHON", "NULL"}: return None - if buf.op is Ops.PARAM and buf.addrspace is AddrSpace.GLOBAL and (dims:=ImageDType.valid_dims(buf.dtype, ctx.target.arch)): - return buf.replace(dtype=(dtypes.imageh if buf.dtype.base == dtypes.half else dtypes.imagef)((*dims[0], 4))) - return None - -def _can_image(buf:UOp, memory:_Memory) -> bool: - # A PARAM slot is image-backed only if every access can be emitted as aligned vec4 pixels. - for (_,membuf,base,_), offsets in memory.items(): - if membuf is not buf: continue - for full_grp in _grouped_offsets(offsets): - while full_grp: - if len(full_grp) < 4 or _offset(base, full_grp[0]).divides(4) is None: return False - full_grp = full_grp[4:] - return True - -def _image_idx(buf:UOp, offset:UOp) -> UOp: - pix = offset // 4 - return buf.index(pix // buf.dtype.shape[1], pix % buf.dtype.shape[1], ptr=True) - -def _coalesced_idx(buf:UOp, ibuf:UOp|None, offset:UOp, length:int) -> UOp: - if ibuf is not None: return _image_idx(ibuf, offset) - return buf._mop(Ops.SHRINK, arg=[(offset, length)]) if length > 1 else buf.index(offset) - -def _fold_lengths(buf:UOp, ctx:Renderer|None, use_image:bool) -> tuple[list[int], bool]: - if use_image: return [4], True - if ctx is not None and ctx.target.device == "DSP": return [128,64,32,16,8,4,1], False - if buf.addrspace == AddrSpace.REG or buf.dtype.base not in (dtypes.float, dtypes.half, *dtypes.fp8s): return [1], True - if ctx is not None and ctx.supports_float4: - return ([8,4,2,1] if buf.dtype.base == dtypes.half and getenv("ALLOW_HALF8") else [4,2,1]), True - return [1], True - -def _choose_length(lengths:list[int], group:list[int], offset:UOp, must_divide:bool) -> int: - return next(l for l in lengths if l <= len(group) and (not must_divide or offset.divides(l) is not None)) - -def _image_buffers(memory:_Memory, ctx:Renderer|None) -> dict[UOp, UOp]: - image_bufs = {} - for _,buf,_,_ in memory: - if buf in image_bufs: continue - if (ibuf:=_image_buf(buf, ctx)) is not None and _can_image(buf, memory): image_bufs[buf] = ibuf - return image_bufs - def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp: if getenv("DMC"): return sink @@ -86,33 +26,79 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp: if u.src[0].op is not Ops.INDEX or len(u.src[0].src) != 2: continue buf, idx_u = u.src[0].src if buf.addrspace == AddrSpace.REG: continue - idx, valid = _valid_idx(idx_u) - root_src, arg = _base_offset(idx) + idx: Any = idx_u.src[1] if idx_u.op is Ops.WHERE and idx_u.src[2].arg is Invalid else idx_u + valid: Any = idx_u.src[0] if idx_u.op is Ops.WHERE and idx_u.src[2].arg is Invalid else None + if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg + elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg + elif idx.op is Ops.CONST and idx.arg is Invalid: root_src, arg = "INVALID", 0 + elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg + else: root_src, arg = idx, 0 memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u) - image_bufs = _image_buffers(memory, ctx) + image_bufs = {} + if IMAGE and ctx is not None and ctx.target.device in {"QCOM", "CL", "PYTHON", "NULL"}: + for _,buf,_,_ in memory: + if buf in image_bufs: continue + ibuf = buf if isinstance(buf.dtype, ImageDType) else None + if ibuf is None and buf.op is Ops.PARAM and buf.addrspace is AddrSpace.GLOBAL and (dims:=ImageDType.valid_dims(buf.dtype, ctx.target.arch)): + ibuf = buf.replace(dtype=(dtypes.imageh if buf.dtype.base == dtypes.half else dtypes.imagef)((*dims[0], 4))) + if ibuf is None: continue + can_image = True + for (_,membuf,base,_), offsets in memory.items(): + if membuf is not buf: continue + for full_grp in [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]: + while full_grp: + offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(dtypes.int, full_grp[0]) + if len(full_grp) < 4 or offset.divides(4) is None: + can_image = False + break + full_grp = full_grp[4:] + if not can_image: break + if not can_image: break + if can_image: image_bufs[buf] = ibuf # build replacements replacements = {} for (op,buf,base,valid),offsets in memory.items(): if isinstance(buf.dtype, ImageDType) and buf not in image_bufs: continue - lengths, must_divide = _fold_lengths(buf, ctx, buf in image_bufs) - for full_grp in _grouped_offsets(offsets): - while full_grp: - offset = _offset(base, full_grp[0]) - length = _choose_length(lengths, full_grp, offset, must_divide) + # allowed lengths (copied in) + lengths = [] + must_divide = True + if buf in image_bufs: + lengths = [4] + elif ctx is not None and ctx.target.device == "DSP": + lengths = [128,64,32,16,8,4] + must_divide = False + elif buf.dtype.base not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not isinstance(buf.dtype, ImageDType): + pass + elif buf.addrspace == AddrSpace.REG: + pass + elif ctx is not None and ctx.supports_float4: + # TODO: a better way to get this than ctx + lengths = [8,4,2] if buf.dtype.base == dtypes.half and getenv("ALLOW_HALF8") else [4,2] + if buf not in image_bufs: lengths.append(1) # worst case, it's not folded + # do the grouping + grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])] + for full_grp in grouped_offsets: + while len(full_grp): + offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(dtypes.int, full_grp[0]) + length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.divides(l) is not None)][0] grp = full_grp[:length] ibuf = image_bufs.get(buf) if length == 4 else None - idx = _coalesced_idx(buf, ibuf, offset, length) + if ibuf is not None: + pix = offset // 4 + idx = ibuf.index(pix // ibuf.dtype.shape[1], pix % ibuf.dtype.shape[1], ptr=True) + else: + idx = buf._mop(Ops.SHRINK, arg=[(offset, len(grp))]) if len(grp) > 1 else buf.index(offset) if op == Ops.STORE: datas = [] - for g in grp: + for i,g in enumerate(grp): assert len(offsets[g]) == 1, f"attempting multiple stores: {len(offsets[g])}" datas.append(offsets[g][0].src[1]) data = UOp.vectorize(*datas) if len(datas) > 1 else datas[0] if ibuf is not None and ibuf.dtype.itemsize == 2: data = pm_imageh_store.rewrite(data) store = idx.store(data, valid) if valid is not None else idx.store(data) - for g in grp: replacements[offsets[g][0]] = store + for i,g in enumerate(grp): replacements[offsets[g][0]] = store else: ld = idx.load(idx.vconst_like(0), valid) if valid is not None else idx.load() if ibuf is not None and buf.dtype.base != dtypes.float: ld = ld.cast(buf.dtype.base)