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

format

Browse Source
This commit is contained in:
Crutcher Dunnavant
2023-04-07 19:12:05 -07:00
parent 34464b40bb
commit 00d382dbf7
2 changed files with 31 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
python/src/_py_mscclpp.cpp
View File

@@ -128,18 +128,18 @@ struct _Comm {
};
struct _P2PHandle {
struct mscclppRegisteredMemoryP2P _rmP2P;
struct mscclppIbMr _ibmr;
struct mscclppRegisteredMemoryP2P _rmP2P;
struct mscclppIbMr _ibmr;
_P2PHandle() : _rmP2P({0}), _ibmr({0}) {}
_P2PHandle() : _rmP2P({0}), _ibmr({0}) {}
_P2PHandle(const mscclppRegisteredMemoryP2P &p2p): _ibmr({0}) {
_rmP2P = p2p;
if (_rmP2P.IbMr != nullptr) {
_ibmr = *_rmP2P.IbMr;
_rmP2P.IbMr = &_ibmr;
}
_P2PHandle(const mscclppRegisteredMemoryP2P& p2p) : _ibmr({0}) {
_rmP2P = p2p;
if (_rmP2P.IbMr != nullptr) {
_ibmr = *_rmP2P.IbMr;
_rmP2P.IbMr = &_ibmr;
}
}
};
nb::callable _log_callback;
@@ -261,7 +261,9 @@ NB_MODULE(_py_mscclpp, m) {
.def_ro("world_size", &_Comm::_world_size)
.def(
"register_buffer",
[](_Comm& comm, uint64_t local_buff, uint64_t buff_size) -> std::vector<_P2PHandle> {
[](_Comm& comm,
uint64_t local_buff,
uint64_t buff_size) -> std::vector<_P2PHandle> {
comm.check_open();
mscclppRegisteredMemory regMem;
checkResult(
@@ -273,8 +275,8 @@ NB_MODULE(_py_mscclpp, m) {
"Registering buffer failed");
std::vector<_P2PHandle> handles;
for (const auto &p2p : regMem.p2p) {
handles.push_back(_P2PHandle(p2p));
for (const auto& p2p : regMem.p2p) {
handles.push_back(_P2PHandle(p2p));
}
return handles;
},