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