diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 1615ea619a..e5efff7f18 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -122,11 +122,13 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: supported_ops = tuple(ren.code_for_op.keys()) pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops) sink = graph_rewrite(sink, pm_decomp, name="early decompositions") - sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes") # do memory coalesing (late) sink = memory_coalesing(sink, ren) + # do dtype decomps + sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes") + # 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 5b5b604689..e59cbb22a1 100644 --- a/tinygrad/codegen/late/coalese.py +++ b/tinygrad/codegen/late/coalese.py @@ -7,7 +7,7 @@ from tinygrad.helpers import getenv, IMAGE from tinygrad.renderer import Renderer from tinygrad.codegen.late.devectorizer import image_valid_dims, _drop_valid_stmts, uop_given_valid -def transform_to_image(ctx, buf:UOp, x:UOp, valid:UOp|None=None) -> UOp: +def transform_to_image(ctx, buf:UOp, x:UOp, valid:UOp|None=None) -> UOp|None: # search for dims that drop the most valid statements best_drop, cands = -1, [] for ch, cw in image_valid_dims(buf.dtype.base, buf.max_numel(), ctx.target.arch):