diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 6af30fb721..fce8fbf756 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -4,6 +4,7 @@ import itertools from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp +from tinygrad.uop.ops import shape_to_shape_arg from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, spec_tensor, spec_program from tinygrad.renderer import Renderer, Estimates @@ -58,9 +59,13 @@ pm_no_weakints = PatternMatcher([ ]) pm_fix_image_shrink = PatternMatcher([ + # proper image index (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]//4) if len(img.shape) == 3 and out.shape == (4,) else None), -]) + # demote this to non image + (UPat(Ops.INDEX, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.PARAM, name="img"), UPat())), UPat(name="idx")), name="out"), + lambda img,out,idx: img.replace(src=(shape_to_shape_arg((img.max_numel(),)),), dtype=dtypes.float).index(idx) if out.shape == () else None) +])+pm_mops def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST") @@ -134,7 +139,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # image fixup if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON", "NULL"}: - sink = graph_rewrite(sink, pm_mops+pm_fix_image_shrink, name="fix image shrink") + sink = graph_rewrite(sink, pm_fix_image_shrink, name="fix image shrink") # instruction selection decompositions pm_decomp = pm_decomp+\