Add Python get_include() (#141)

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.

CMakeLists.txt

@@ -83,11 +83,13 @@ set_target_properties(mscclpp_static PROPERTIES VERSION ${MSCCLPP_VERSION} SOVER
 
 add_subdirectory(include)
 add_subdirectory(src)
-install(TARGETS mscclpp mscclpp_static
-        LIBRARY DESTINATION lib
-        ARCHIVE DESTINATION lib)
 
-install(DIRECTORY ${PROJECT_SOURCE_DIR}/include/ DESTINATION include FILES_MATCHING PATTERN "*.hpp")
+install(TARGETS mscclpp_obj
+    FILE_SET HEADERS DESTINATION include)
+install(TARGETS mscclpp
+    LIBRARY DESTINATION lib)
+install(TARGETS mscclpp_static
+    ARCHIVE DESTINATION lib)
 
 # Tests
 if (BUILD_TESTS)

include/CMakeLists.txt

@@ -3,6 +3,3 @@
 
 file(GLOB_RECURSE HEADERS CONFIGURE_DEPENDS *.hpp)
 target_sources(mscclpp_obj PUBLIC FILE_SET HEADERS FILES ${HEADERS})
-if(NOT SKBUILD)
-    install(TARGETS mscclpp FILE_SET HEADERS)
-endif()

pyproject.toml

@@ -1,3 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
 [build-system]
 requires = ["scikit-build-core"]
 build-backend = "scikit_build_core.build"
@@ -9,6 +12,8 @@ version = "0.2.0"
 [tool.scikit-build]
 cmake.minimum-version = "3.25.0"
 build-dir = "build/{wheel_tag}"
+wheel.packages = ["python/mscclpp"]
+wheel.install-dir = "mscclpp"
 
 [tool.scikit-build.cmake.define]
 BUILD_PYTHON_BINDINGS = "ON"

python/CMakeLists.txt

@@ -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 .)

python/core_py.cpp

@@ -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);

python/mscclpp/__init__.py

@@ -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")