forked from tinygrad/tinygrad
simpler
This commit is contained in:
@@ -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()
|
||||
@@ -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+\
|
||||
|
||||
@@ -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<imageh>(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)
|
||||
|
||||
Reference in New Issue
Block a user