mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-05-12 01:10:22 +00:00
update python binding (#136)
update pythons binding for `device_handle`
This commit is contained in:
23
python/sm_channel_py.cpp
Normal file
23
python/sm_channel_py.cpp
Normal 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"));
|
||||
};
|
||||
Reference in New Issue
Block a user