microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-19 02:06:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

merge main

Browse Source
This commit is contained in:
Binyang Li
2026-07-13 22:01:36 +00:00
parent 359e681e8e 9b42ad3062
commit bc4374e650
102 changed files with 4323 additions and 576 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/csrc/core_py.cpp
View File

@@ -56,6 +56,7 @@ void register_core(nb::module_& m) {
.def("get_rank", &Bootstrap::getRank)
.def("get_n_ranks", &Bootstrap::getNranks)
.def("get_n_ranks_per_node", &Bootstrap::getNranksPerNode)
.def("get_n_ranks_per_ipc_domain", &Bootstrap::getNranksPerIpcDomain)
.def(
"send",
[](Bootstrap* self, uintptr_t ptr, size_t size, int peer, int tag) {

7
python/csrc/gpu_utils_py.cpp
View File

@@ -114,8 +114,13 @@ static nb::capsule toDlpack(GpuBuffer<char> buffer, std::string dataType, std::v
void register_gpu_utils(nb::module_& m) {
m.def("is_nvls_supported", &isNvlsSupported);
nb::enum_<GpuBufferGranularity>(m, "CppGpuBufferGranularity")
.value("MultiCastMinimum", GpuBufferGranularity::MultiCastMinimum)
.value("MultiCastRecommended", GpuBufferGranularity::MultiCastRecommended);
nb::class_<GpuBuffer<char>>(m, "CppRawGpuBuffer")
.def(nb::init<size_t>(), nb::arg("nelems"))
.def(nb::init<size_t, GpuBufferGranularity>(), nb::arg("nelems"),
nb::arg("granularity") = GpuBufferGranularity::MultiCastMinimum)
.def("nelems", &GpuBuffer<char>::nelems)
.def("bytes", &GpuBuffer<char>::bytes)
.def("data", [](GpuBuffer<char>& self) { return reinterpret_cast<uintptr_t>(self.data()); })