diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 0d6f4cc670..c823ef4777 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -57,6 +57,11 @@ pm_no_weakints = PatternMatcher([ (UPat(GroupOp.All, dtype=dtypes.weakint, name="x"), lambda x: x.replace(dtype=dtypes.int)) ]) +pm_fix_image_shrink = PatternMatcher([ + (UPat(Ops.SHRINK, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.PARAM, name="img"), UPat())), UPat(name="idx"), UPat()), name="out"), + lambda img,out,idx: img.reshape(-1, 4).index(idx.src[0]) if len(img.shape) == 3 else None), +]) + def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST") if DEBUG >= 5: print(pyrender(ast)) @@ -103,10 +108,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # add loads and remove invalids sink = graph_rewrite(sink, pm_add_loads+pm_remove_invalid, name="** add loads (code)") - # create image buffers - if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON", "NULL"}: - sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True, ctx=ren.target.arch) - # devectorize sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding+correct_load_store+load_store_indexing, ctx=ren, name="devectorize") @@ -124,9 +125,16 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = graph_rewrite(sink, pm_decomp, name="early decompositions") sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes") + # create image buffers + if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON", "NULL"}: + sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True, ctx=ren.target.arch) + # do memory coalesing (late) sink = memory_coalesing(sink, ren) + # image fixup + sink = graph_rewrite(sink, pm_mops+pm_fix_image_shrink, name="fix image shrink") + # instruction selection decompositions pm_decomp = pm_decomp+\ get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\ diff --git a/tinygrad/codegen/late/coalese.py b/tinygrad/codegen/late/coalese.py index fb939dca29..78494a76a0 100644 --- a/tinygrad/codegen/late/coalese.py +++ b/tinygrad/codegen/late/coalese.py @@ -13,7 +13,7 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp: memory: defaultdict[tuple[Ops, UOp, Any, Any], dict[int, list[UOp]]] = defaultdict(dict) for u in sink.toposort(): # TODO: this should handle images too, it's just memory coalesing - if u.op in {Ops.LOAD, Ops.STORE} and not isinstance(u.src[0].src[0].dtype, ImageDType): + if u.op in {Ops.LOAD, Ops.STORE}: assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalesing does not support gated loads/stores" if u.src[0].op is not Ops.INDEX: continue buf, idx_u = u.src[0].src