gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FRONTEND] Disable P2P (#312)

Browse Source
This commit is contained in:
Philippe Tillet
2021-09-27 21:18:27 -07:00
committed by GitHub
parent 2c287544cb
commit bfcfad7abe
1 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/triton/code_gen.py
View File

@@ -574,15 +574,15 @@ class Kernel:
device = torch.device('cuda', torch.cuda.current_device())
device_idx = device.index
if len(set(device_ids)) != 1 or device_ids[0] != device_idx:
# try to enable P2P communication
for arg_idx, dst_idx in zip(tensor_idxs, device_ids):
if dst_idx != device_idx:
try:
_triton.runtime.enable_peer_access(self.backend, wargs[arg_idx].data_ptr())
except RuntimeError as e:
raise RuntimeError("Cannot enable P2P access from device {} to device {}: {}"
.format(device_idx, dst_idx, str(e)))
# if len(set(device_ids)) != 1 or device_ids[0] != device_idx:
# # try to enable P2P communication
# for arg_idx, dst_idx in zip(tensor_idxs, device_ids):
# if dst_idx != device_idx:
# try:
# _triton.runtime.enable_peer_access(self.backend, wargs[arg_idx].data_ptr())
# except RuntimeError as e:
# raise RuntimeError("Cannot enable P2P access from device {} to device {}: {}"
# .format(device_idx, dst_idx, str(e)))
# enqueue kernel on the current device
torch.cuda.set_device(device_idx)