tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 10:56:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

coalese

Browse Source
This commit is contained in:
geohot
2026-06-23 09:04:23 -07:00
parent 87e12fad81
commit 5238f304c7
2 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/__init__.py
+1
View File
@@ -114,6 +114,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
# ctx=ren, name="devectorize")
sink = graph_rewrite(sink, unbroadcast, name="*** unbroadcast")
sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
sink = graph_rewrite(sink, symbolic, name="pre memory coalese")
# memory coalesing
sink = memory_coalesing(sink)
tinygrad/codegen/codegen2.py
+1 -1
View File
@@ -116,7 +116,7 @@ def memory_coalesing(sink:UOp):
# collect
memory: defaultdict[tuple[UOp, UOp, UOp], dict[int, list[UOp]]] = defaultdict(dict)
for u in sink.toposort():
if u.op in {Ops.LOAD, Ops.STORE}:
if u.op in {Ops.LOAD, Ops.STORE} and u.src[0].addrspace != AddrSpace.REG:
assert u.src[0].op is Ops.INDEX
buf,idx_u = u.src[0].src
idx: Any = idx_u.get_idx()