github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

All MultiGraphRunners can graph transfers (#10301)

Browse Source
This commit is contained in:
uuuvn
2025-05-14 19:23:02 +05:00
committed by GitHub
parent e00679dc92
commit dd816d0237
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tinygrad/engine/jit.py
View File

@@ -48,8 +48,8 @@ def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer]
can_be_graphed = ji_graph_dev.graph is not None
case BufferXfer():
ji_graph_dev = Device[unwrap(ji.bufs[0]).device]
# Whitelist of devices that support graphing BufferXfers
can_be_graphed = ji_graph_dev.graph is not None and ji_graph_dev.device.split(":", 1)[0] in {"CUDA", "NV", "AMD", "NULL"}
# All *Multi*GraphRunner support graphing BufferXfers
can_be_graphed = ji_graph_dev.graph is not None and issubclass(graph_class(ji_graph_dev), MultiGraphRunner)
case ViewOp(): continue # ViewOps are just ignored
case _: can_be_graphed = False # Everything else is not graphed and flushes existing graph if it's being constructed