diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 77018c1a59..2b131ecb7d 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -24,7 +24,7 @@ from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_s from tinygrad.schedule.rangeify import pm_mops from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite -from tinygrad.codegen.late.coalesce import memory_coalesing, pm_simplify_add_image +from tinygrad.codegen.late.coalesce import memory_coalescing, pm_simplify_add_image from tinygrad.helpers import all_same, flatten, argsort, partition from tinygrad.uop.ops import _align_left, _broadcast_shape, identity_element from tinygrad.schedule.rangeify import BufferizeOpts @@ -315,11 +315,11 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # simplify indexing sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing") - # some coalesing misses without this + # some coalescing misses without this sink = graph_rewrite(sink, sym, name="early symbolic") - # do memory coalesing (late) - sink = memory_coalesing(sink, ren) + # do memory coalescing (late) + sink = memory_coalescing(sink, ren) sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True) # extra symbolic before decomp. crashes without this? diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index eeee6a4b1c..24429c44d6 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -97,16 +97,16 @@ pm_simplify_add_image = PatternMatcher([ (UPat.var("x", dtype=dtypes.float).cast(dtypes.half).cast(dtypes.float), lambda x: x), ]) -def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp: +def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: if getenv("DMC"): return sink # collect memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp], dict[int, list[UOp]]] = defaultdict(dict) for u in sink.toposort(): - # TODO: this should handle images too, it's just memory coalesing + # TODO: this should handle images too, it's just memory coalescing 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" - assert u.src[0].op is Ops.INDEX, f"memory coalesing should be on INDEX, not {u.src[0].op}" + assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores" + assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}" buf, idx_u = u.src[0].src if buf.addrspace == AddrSpace.REG: continue idx, valid = idx_u.get_idx(), idx_u.get_valid() @@ -162,4 +162,4 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp: full_grp = full_grp[length:] # apply - return sink.substitute(replacements, name="memory coalesing") + return sink.substitute(replacements, name="memory coalescing")