microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-05-13 01:36:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Get rid of comm.setup()

Browse Source
This commit is contained in:
Olli Saarikivi
2023-08-31 17:45:58 +00:00
committed by Saeed Maleki
parent 0863e862f5
commit 8cb63a7d1a
23 changed files with 253 additions and 352 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/examples/bootstrap.py
View File

@@ -40,12 +40,11 @@ def setup_connections(comm, rank, world_size, element_size, proxy_service):
for r in range(world_size):
if r == rank:
continue
conn = comm.connect_on_setup(r, 0, mscclpp.Transport.CudaIpc)
conn = comm.connect(r, 0, mscclpp.Transport.CudaIpc)
connections.append(conn)
comm.send_memory_on_setup(reg_mem, r, 0)
remote_mem = comm.recv_memory_on_setup(r, 0)
comm.send_memory(reg_mem, r, 0)
remote_mem = comm.recv_memory(r, 0)
remote_memories.append(remote_mem)
comm.setup()
connections = [conn.get() for conn in connections]

8
python/examples/send_recv.py
View File

@@ -35,15 +35,13 @@ def main(args):
size = elements * memory.itemsize
my_reg_mem = comm.register_memory(ptr, size, mscclpp.Transport.IB0)
conn = comm.connect_on_setup((rank + 1) % 2, 0, mscclpp.Transport.IB0)
conn = comm.connect((rank + 1) % 2, 0, mscclpp.Transport.IB0)
other_reg_mem = None
if rank == 0:
other_reg_mem = comm.recv_memory_on_setup((rank + 1) % 2, 0)
other_reg_mem = comm.recv_memory((rank + 1) % 2, 0)
else:
comm.send_memory_on_setup(my_reg_mem, (rank + 1) % 2, 0)
comm.setup()
comm.send_memory(my_reg_mem, (rank + 1) % 2, 0)
if rank == 0:
other_reg_mem = other_reg_mem.get()