mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-05-16 10:59:59 +00:00
5acac93dbc
Provides two integration ways for MSCCL++ DSL. 1. Integrate with customized communication group 2. Integrate with NCCL API Introduce new Python APIs to make it work: ```python mscclpp.compile # compile dsl to json based execution plan mscclpp.ExecutionPlanRegistry.register_plan(plan) # register the compiled plan to executionPlanRegistery mscclpp.ExecutionPlanRegistry.set_selector(selector) # set the selector, the selector will return the best execution plan based on collection, message size, world size.... ``` Fix #556 --------- Co-authored-by: Caio Rocha <caiorocha@microsoft.com> Co-authored-by: Changho Hwang <changhohwang@microsoft.com>
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"));
|
|
}
|