mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-05-03 13:11:27 +00:00
e21513791a
Rename nanobind-exposed C++ types to Cpp* Replace MSCCLPP_EXECUTION_PLAN_DIR / MSCCLPP_NATIVE_CACHE_DIR with MSCCLPP_CACHE_DIR across C++ and Python.
34 lines
1.5 KiB
C++
34 lines
1.5 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// Licensed under the MIT License.
|
|
|
|
#include <nanobind/nanobind.h>
|
|
#include <nanobind/stl/function.h>
|
|
#include <nanobind/stl/shared_ptr.h>
|
|
#include <nanobind/stl/unordered_map.h>
|
|
#include <nanobind/stl/vector.h>
|
|
|
|
#include <mscclpp/algorithm.hpp>
|
|
#include <mscclpp/ext/collectives/algorithm_collection_builder.hpp>
|
|
|
|
namespace nb = nanobind;
|
|
using namespace mscclpp;
|
|
using namespace mscclpp::collective;
|
|
|
|
void register_algorithm_collection_builder(nb::module_& m) {
|
|
nb::class_<AlgorithmCollectionBuilder>(m, "CppAlgorithmCollectionBuilder")
|
|
.def_static("get_instance", &AlgorithmCollectionBuilder::getInstance)
|
|
.def("add_algorithm_builder", &AlgorithmCollectionBuilder::addAlgorithmBuilder, nb::arg("builder"))
|
|
.def(
|
|
"add_dsl_algorithm_builder",
|
|
[](AlgorithmCollectionBuilder& self, std::shared_ptr<DslAlgorithm> algorithm) {
|
|
self.addAlgorithmBuilder(algorithm);
|
|
},
|
|
nb::arg("algorithm"))
|
|
.def("set_algorithm_selector", &AlgorithmCollectionBuilder::setAlgorithmSelector, nb::arg("selector"))
|
|
.def("set_fallback_algorithm_selector", &AlgorithmCollectionBuilder::setFallbackAlgorithmSelector,
|
|
nb::arg("selector"))
|
|
.def("build", &AlgorithmCollectionBuilder::build)
|
|
.def("build_default_algorithms", &AlgorithmCollectionBuilder::buildDefaultAlgorithms, nb::arg("scratch_buffer"),
|
|
nb::arg("scratch_buffer_size"), nb::arg("rank"))
|
|
.def_static("reset", &AlgorithmCollectionBuilder::reset);
|
|
} |