mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-05-24 23:06:17 +00:00
09219c1f6a
* Python cannot distinguish `Communicator::connect(const Endpoint&, ...)` from `Communicator::connect(const EndpointConfig&, ...)`. Temporarily removed the former one. * A few other fixes in Python bindings.
38 lines
1.5 KiB
C++
38 lines
1.5 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// Licensed under the MIT license.
|
|
|
|
#include <nanobind/nanobind.h>
|
|
#include <nanobind/operators.h>
|
|
#include <nanobind/stl/array.h>
|
|
#include <nanobind/stl/shared_ptr.h>
|
|
#include <nanobind/stl/string.h>
|
|
#include <nanobind/stl/vector.h>
|
|
|
|
#include <mscclpp/core.hpp>
|
|
#include <mscclpp/switch_channel.hpp>
|
|
|
|
namespace nb = nanobind;
|
|
using namespace mscclpp;
|
|
|
|
void register_nvls(nb::module_& m) {
|
|
nb::class_<SwitchChannel>(m, "SwitchChannel")
|
|
.def("get_device_ptr", [](SwitchChannel* self) { return (uintptr_t)self->getDevicePtr(); })
|
|
.def("device_handle", &SwitchChannel::deviceHandle);
|
|
|
|
nb::class_<SwitchChannel::DeviceHandle>(m, "DeviceHandle")
|
|
.def(nb::init<>())
|
|
.def_rw("device_ptr", &SwitchChannel::DeviceHandle::devicePtr)
|
|
.def_rw("mc_ptr", &SwitchChannel::DeviceHandle::mcPtr)
|
|
.def_rw("size", &SwitchChannel::DeviceHandle::bufferSize)
|
|
.def_prop_ro("raw", [](const SwitchChannel::DeviceHandle& self) -> nb::bytes {
|
|
return nb::bytes(reinterpret_cast<const char*>(&self), sizeof(self));
|
|
});
|
|
|
|
nb::class_<NvlsConnection>(m, "NvlsConnection")
|
|
.def("bind_allocated_memory", &NvlsConnection::bindAllocatedMemory, nb::arg("device_ptr"), nb::arg("size"))
|
|
.def("get_multicast_min_granularity", &NvlsConnection::getMultiCastMinGranularity);
|
|
|
|
m.def("connect_nvls_collective", &connectNvlsCollective, nb::arg("communicator"), nb::arg("all_ranks"),
|
|
nb::arg("buffer_size"));
|
|
}
|