mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-05-11 08:50:21 +00:00
35 lines
1.6 KiB
C++
35 lines
1.6 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/string.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("flag_buffer"), nb::arg("flag_buffer_size"), nb::arg("rank"))
|
|
.def_static("reset", &AlgorithmCollectionBuilder::reset);
|
|
} |