update python binding (#136)

update pythons binding for `device_handle`
This commit is contained in:
Binyang2014
2023-07-24 11:00:33 +08:00
committed by GitHub
parent a58e2e9623
commit 9a488f0da2
5 changed files with 32 additions and 4 deletions

23
python/sm_channel_py.cpp Normal file
View File

@@ -0,0 +1,23 @@
// 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(nb::init<std::shared_ptr<SmDevice2DeviceSemaphore>, RegisteredMemory, void*, void*>(), nb::arg("semaphore"),
nb::arg("dst"), nb::arg("src"), nb::arg("getPacketBuffer"))
.def("device_handle", &SmChannel::deviceHandle);
nb::class_<SmChannel::DeviceHandle>(smChannel, "DeviceHandle");
m.def("device_handle", &deviceHandle<SmChannel>, nb::arg("smChannel"));
};