nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

correct more spelling of coalesce (#17107)

Browse Source
This commit is contained in:
chenyu
2026-07-20 21:55:37 -04:00
committed by GitHub
parent f3a5337825
commit 2864036e8e
2 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/__init__.py
+4 -4
View File
@@ -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?
tinygrad/codegen/late/coalesce.py
+5 -5
View File
@@ -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")