Files
mscclpp/python/sm_channel_py.cpp
Saeed Maleki 8d1b984bed Change device handle interfaces & others (#142)
* Changed device handle interfaces
* Changed proxy service interfaces
* Move device code into separate files
* Fixed FIFO polling issues
* Add configuration arguments in several interface functions

---------

Co-authored-by: Changho Hwang <changhohwang@microsoft.com>
Co-authored-by: Binyang Li <binyli@microsoft.com>
Co-authored-by: root <root@a100-saemal0.qxveptpukjsuthqvv514inp03c.gx.internal.cloudapp.net>
2023-08-16 20:00:56 +08:00

36 lines
1.4 KiB
C++

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
#include <nanobind/nanobind.h>
#include <nanobind/stl/shared_ptr.h>
#include <nanobind/stl/string.h>
#include <mscclpp/sm_channel.hpp>
namespace nb = nanobind;
using namespace mscclpp;
void register_sm_channel(nb::module_& m) {
nb::class_<SmChannel> smChannel(m, "SmChannel");
smChannel
.def("__init__",
[](SmChannel* smChannel, std::shared_ptr<SmDevice2DeviceSemaphore> semaphore, RegisteredMemory dst,
uintptr_t src) { new (smChannel) SmChannel(semaphore, dst, (void*)src); })
.def("__init__",
[](SmChannel* smChannel, std::shared_ptr<SmDevice2DeviceSemaphore> semaphore, RegisteredMemory dst,
uintptr_t src, uintptr_t get_packet_buffer) {
new (smChannel) SmChannel(semaphore, dst, (void*)src, (void*)get_packet_buffer);
})
.def("device_handle", &SmChannel::deviceHandle);
nb::class_<SmChannel::DeviceHandle>(m, "SmChannelDeviceHandle")
.def(nb::init<>())
.def_rw("semaphore_", &SmChannel::DeviceHandle::semaphore_)
.def_rw("src_", &SmChannel::DeviceHandle::src_)
.def_rw("dst_", &SmChannel::DeviceHandle::dst_)
.def_rw("getPacketBuffer_", &SmChannel::DeviceHandle::getPacketBuffer_)
.def_prop_ro("raw", [](const SmChannel::DeviceHandle& self) -> nb::bytes {
return nb::bytes(reinterpret_cast<const char*>(&self), sizeof(self));
});
};