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

Queue remote transfers on source (#10871)

Browse Source 
https://github.com/tinygrad/tinygrad/pull/10601#issuecomment-2985624147

I personally don't see how that is a good standalone pr, but whatever
This commit is contained in:
uuuvn
2025-06-18 16:08:44 -07:00
committed by GitHub
parent a9f3632c4f
commit e7a26211d2
2 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/runtime/graph/remote.py
+2 -2
View File
@@ -22,8 +22,8 @@ class RemoteGraph(GraphRunner):
case BufferXfer():
dest, src = ji.bufs[0:2]
assert dest is not None and src is not None, ji
return Transfer(session=cast(RemoteDevice, Device[dest.device]).session, buffer_num=dest._buf,
ssession=cast(RemoteDevice, Device[src.device]).session, sbuffer_num=src._buf)
return Transfer(session=cast(RemoteDevice, Device[src.device]).session, buffer_num=src._buf,
dsession=cast(RemoteDevice, Device[dest.device]).session, dbuffer_num=dest._buf)
self.graph_num = next(self.devices[0].graph_num)
self.devices[0].q(GraphAlloc(self.graph_num, tuple(_process_ji(ji) for ji in jit_cache), self.map_rawbufs(rawbufs), var_vals))
tinygrad/runtime/ops_remote.py
+7 -7
View File
@@ -56,7 +56,7 @@ class CopyIn(RemoteRequest): buffer_num: int; datahash: str # noqa: E702
class CopyOut(RemoteRequest): buffer_num: int
@dataclass(frozen=True)
class Transfer(RemoteRequest): buffer_num: int; ssession: SessionKey; sbuffer_num: int # noqa: E702
class Transfer(RemoteRequest): buffer_num: int; dsession: SessionKey; dbuffer_num: int # noqa: E702
@dataclass(frozen=True)
class ProgramAlloc(RemoteRequest): name: str; datahash: str # noqa: E702
@@ -190,11 +190,11 @@ class RemoteHandler:
case CopyIn(): session.buffers[c.buffer_num].copyin(memoryview(bytearray(req._h[c.datahash])))
case CopyOut(): session.buffers[c.buffer_num].copyout(memoryview(ret:=bytearray(session.buffers[c.buffer_num].nbytes)))
case Transfer():
ssession, sdev = self.sessions[c.ssession], Device[f"{self.base_device}:{unwrap(c.ssession).idx}"]
dbuf, sbuf = session.buffers[c.buffer_num], ssession.buffers[c.sbuffer_num]
dsession, ddev = self.sessions[c.dsession], Device[f"{self.base_device}:{unwrap(c.dsession).idx}"]
dbuf, sbuf = dsession.buffers[c.dbuffer_num], session.buffers[c.buffer_num]
assert dbuf.nbytes == sbuf.nbytes, f"{dbuf.nbytes} != {sbuf.nbytes}"
assert hasattr(dev.allocator, '_transfer'), f"Device {dev.device} doesn't support transfers"
dev.allocator._transfer(dbuf._buf, sbuf._buf, dbuf.nbytes, dest_dev=dev, src_dev=sdev)
assert hasattr(ddev.allocator, '_transfer'), f"Device {ddev.device} doesn't support transfers"
ddev.allocator._transfer(dbuf._buf, sbuf._buf, dbuf.nbytes, dest_dev=ddev, src_dev=dev)
case ProgramAlloc():
lib = dev.compiler.compile_cached(req._h[c.datahash].decode())
session.programs[(c.name, c.datahash)] = dev.runtime(c.name, lib)
@@ -217,7 +217,7 @@ class RemoteHandler:
return ExecItem(CompiledRunner(ps, precompiled=b'', prg=prg), [self.sessions[gi.session].buffers[buf] for buf in gi.bufs],
fixedvars=gi.fixedvars)
case Transfer():
dbuf, sbuf = self.sessions[unwrap(gi.session)].buffers[gi.buffer_num], self.sessions[gi.ssession].buffers[gi.sbuffer_num]
dbuf, sbuf = self.sessions[gi.dsession].buffers[gi.dbuffer_num], self.sessions[unwrap(gi.session)].buffers[gi.buffer_num]
assert dbuf.nbytes == sbuf.nbytes, f"{dbuf.nbytes} != {sbuf.nbytes}"
return ExecItem(BufferXfer(dbuf.nbytes, dbuf.device, sbuf.device), [dbuf, sbuf])
assert c.graph_num not in session.graphs, f"graph {c.graph_num} already allocated"
@@ -257,7 +257,7 @@ class RemoteAllocator(Allocator['RemoteDevice']):
dest[:] = resp
def _transfer(self, dest, src, sz, src_dev, dest_dev):
if dest_dev.properties.transfer_supported and src_dev.conn == dest_dev.conn:
dest_dev.q(Transfer(dest, src_dev.session, src))
src_dev.q(Transfer(src, dest_dev.session, dest))
else:
src_dev.allocator._copyout(tmp:=memoryview(bytearray(sz)), src)
dest_dev.allocator._copyin(dest, tmp)