microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-06-07 00:05:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add Python get_include() (#141)

Browse Source 
Introduces a mscclpp.get_include() in the Python module.
The extension module is now named _mscclpp so that we can have
Python code in the mscclpp module.
Also does some miscellaneous cleanup.
This commit is contained in:
Olli Saarikivi
2023-07-25 10:23:16 -07:00
committed by GitHub
parent 9a488f0da2
commit 4865b2017b
6 changed files with 26 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/CMakeLists.txt
View File

@@ -6,11 +6,9 @@ include(FetchContent)
FetchContent_Declare(nanobind GIT_REPOSITORY https://github.com/wjakob/nanobind.git GIT_TAG v1.4.0)
FetchContent_MakeAvailable(nanobind)
nanobind_add_module(mscclpp_py core_py.cpp error_py.cpp proxy_channel_py.cpp fifo_py.cpp semaphore_py.cpp
config_py.cpp utils_py.cpp sm_channel_py.cpp)
set_target_properties(mscclpp_py PROPERTIES OUTPUT_NAME mscclpp)
file(GLOB_RECURSE SOURCES CONFIGURE_DEPENDS *.cpp)
nanobind_add_module(mscclpp_py ${SOURCES})
set_target_properties(mscclpp_py PROPERTIES OUTPUT_NAME _mscclpp)
target_link_libraries(mscclpp_py PRIVATE mscclpp_static)
target_include_directories(mscclpp_py PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
if (SKBUILD)
install(TARGETS mscclpp_py LIBRARY DESTINATION ${SKBUILD_PLATLIB_DIR})
endif()
install(TARGETS mscclpp_py LIBRARY DESTINATION .)

2
python/core_py.cpp
View File

@@ -143,7 +143,7 @@ void register_core(nb::module_& m) {
.def("setup", &Communicator::setup);
}
NB_MODULE(mscclpp, m) {
NB_MODULE(_mscclpp, m) {
register_error(m);
register_proxy_channel(m);
register_sm_channel(m);

10
python/mscclpp/__init__.py Normal file
View File

@@ -0,0 +1,10 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
from ._mscclpp import *
import os as _os
def get_include():
"""Return the directory that contains the MSCCL++ headers."""
return _os.path.join(_os.path.dirname(__file__), "include")