WIP

Author: Binyang2014

apps/nccl/src/nccl.cu

@@ -326,12 +326,8 @@ NCCL_API ncclResult_t ncclCommInitRank(ncclComm_t* comm, int nranks, ncclUniqueI
   commPtr->nRanksPerNode = mscclppComm->bootstrap()->getNranksPerNode();
   commPtr->worldSize = mscclppComm->bootstrap()->getNranks();
   mscclpp::AlgorithmCollectionBuilder::getInstance()->setFallbackAlgorithmSelector(algoSelector);
-  commPtr->algorithmCollection =
-      mscclpp::AlgorithmCollectionBuilder::getInstance()->buildCollectionWithDefaultNativeAlgorithms(
-          reinterpret_cast<uintptr_t>(commPtr->scratchBuffer_.get()), commPtr->scratchBufferSize_);
-  auto dslAlgoCollection =
-      mscclpp::AlgorithmCollectionBuilder::getInstance()->buildCollectionWithDefaultDslAlgorithms(rank);
-  commPtr->algorithmCollection->extend(*dslAlgoCollection);
+  commPtr->algorithmCollection = mscclpp::AlgorithmCollectionBuilder::getInstance()->buildDefaultAlgorithms(
+      reinterpret_cast<uintptr_t>(commPtr->scratchBuffer_.get()), commPtr->scratchBufferSize_, rank);
 
   *comm = commPtr;
 #if defined(ENABLE_NPKIT)

include/mscclpp/algorithm.hpp

@@ -245,7 +245,7 @@ class AlgorithmCollection {
 
   std::unordered_map<std::string, std::shared_ptr<Algorithm>> getAlgorithmsByCollective(
       const std::string& collective) const;
-
+  std::vector<std::shared_ptr<Algorithm>> getAllAlgorithms() const;
   void extend(const AlgorithmCollection& other);
 
  private:
@@ -265,10 +265,8 @@ class AlgorithmCollectionBuilder {
   /// @param builder The algorithm builder.
   void addAlgorithmBuilder(std::shared_ptr<AlgorithmBuilder> builder);
 
-  std::shared_ptr<AlgorithmCollection> buildCollectionWithDefaultNativeAlgorithms(uintptr_t scratchBuffer,
-                                                                                  size_t scratchBufferSize);
-
-  std::shared_ptr<AlgorithmCollection> buildCollectionWithDefaultDslAlgorithms(int rank);
+  std::shared_ptr<AlgorithmCollection> buildDefaultAlgorithms(uintptr_t scratchBuffer, size_t scratchBufferSize,
+                                                              int rank);
 
   /// @brief Set a new algorithm selection function.
   /// @param selector The algorithm selection function.
@@ -291,6 +289,8 @@ class AlgorithmCollectionBuilder {
   AlgoSelectFunc fallbackAlgoSelector_ = nullptr;
 
   static std::shared_ptr<AlgorithmCollectionBuilder> gAlgorithmCollectionBuilder_;
+  std::shared_ptr<AlgorithmCollection> buildDefaultNativeAlgorithms(uintptr_t scratchBuffer, size_t scratchBufferSize);
+  std::shared_ptr<AlgorithmCollection> buildDefaultDslAlgorithms(int rank);
 };
 
 }  // namespace mscclpp

python/csrc/algorithm.cpp

@@ -7,6 +7,7 @@
 #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>
 
@@ -84,12 +85,15 @@ void register_algorithm(nb::module_& m) {
       .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);
 
   nb::class_<AlgorithmCollection>(m, "AlgorithmCollection")
       .def("register_algorithm", &AlgorithmCollection::registerAlgorithm, nb::arg("collective"), nb::arg("algo_name"),
            nb::arg("algorithm"))
-      .def("get_algorithms_by_collective", &AlgorithmCollection::getAlgorithmsByCollective, nb::arg("collective"));
+      .def("get_algorithms_by_collective", &AlgorithmCollection::getAlgorithmsByCollective, nb::arg("collective"))
+      .def("to_list", &AlgorithmCollection::getAllAlgorithms);
 
   nb::class_<CollectiveRequest>(m, "CollectiveRequest")
       .def_ro("world_size", &CollectiveRequest::worldSize)

src/algorithms/algorithm.cc

@@ -91,6 +91,9 @@ void AlgorithmCollection::registerAlgorithm(const std::string collective, const
 
 std::shared_ptr<Algorithm> AlgorithmCollection::selectAlgorithm(const CollectiveRequest& request) {
   std::shared_ptr<Algorithm> algo;
+  if (!algoSelector_ && !fallbackAlgoSelector_) {
+    THROW(ALGO, Error, ErrorCode::InvalidUsage, "No algorithm selector is set in AlgorithmCollection.");
+  }
   if (algoSelector_) {
     algo = algoSelector_(algoMapByCollective_, request);
   }
@@ -108,6 +111,16 @@ void AlgorithmCollection::extend(const AlgorithmCollection& other) {
   }
 }
 
+std::vector<std::shared_ptr<Algorithm>> AlgorithmCollection::getAllAlgorithms() const {
+  std::vector<std::shared_ptr<Algorithm>> allAlgos;
+  for (const auto& [collective, algoMap] : algoMapByCollective_) {
+    for (const auto& [algoName, algorithm] : algoMap) {
+      allAlgos.push_back(algorithm);
+    }
+  }
+  return allAlgos;
+}
+
 std::unordered_map<std::string, std::shared_ptr<Algorithm>> AlgorithmCollection::getAlgorithmsByCollective(
     const std::string& collective) const {
   auto it = algoMapByCollective_.find(collective);
@@ -130,7 +143,16 @@ void AlgorithmCollectionBuilder::addAlgorithmBuilder(std::shared_ptr<AlgorithmBu
   this->algoBuilders_.push_back(builder);
 }
 
-std::shared_ptr<AlgorithmCollection> AlgorithmCollectionBuilder::buildCollectionWithDefaultNativeAlgorithms(
+std::shared_ptr<AlgorithmCollection> AlgorithmCollectionBuilder::buildDefaultAlgorithms(uintptr_t scratchBuffer,
+                                                                                        size_t scratchBufferSize,
+                                                                                        int rank) {
+  auto nativeCollection = buildDefaultNativeAlgorithms(scratchBuffer, scratchBufferSize);
+  auto dslCollection = buildDefaultDslAlgorithms(rank);
+  nativeCollection->extend(*dslCollection);
+  return nativeCollection;
+}
+
+std::shared_ptr<AlgorithmCollection> AlgorithmCollectionBuilder::buildDefaultNativeAlgorithms(
     uintptr_t scratchBuffer, size_t scratchBufferSize) {
   auto collection = std::make_shared<AlgorithmCollection>();
   auto allreduceAllpairPkt =
@@ -163,7 +185,7 @@ std::shared_ptr<AlgorithmCollection> AlgorithmCollectionBuilder::buildCollection
   return collection;
 }
 
-std::shared_ptr<AlgorithmCollection> AlgorithmCollectionBuilder::buildCollectionWithDefaultDslAlgorithms(int rank) {
+std::shared_ptr<AlgorithmCollection> AlgorithmCollectionBuilder::buildDefaultDslAlgorithms(int rank) {
   struct DslAlgoConfig {
     std::string filename;
     std::string collective;
@@ -188,14 +210,14 @@ std::shared_ptr<AlgorithmCollection> AlgorithmCollectionBuilder::buildCollection
 
   std::string planDir = mscclpp::env()->executionPlanDir;
   if (!std::filesystem::exists(planDir)) {
-    INFO(EXEC, "Plan directory does not exist: ", planDir);
+    INFO(ALGO, "Plan directory does not exist: ", planDir);
     return collection;
   }
   for (const auto& config : defaultAlgoConfigs) {
     std::string planPath = planDir + "/" + config.filename;
-    INFO(EXEC, "Loading plan: ", planPath);
+    INFO(ALGO, "Loading plan: ", planPath);
     if (!std::filesystem::exists(planPath)) {
-      INFO(EXEC, "Plan file does not exist: ", planPath);
+      INFO(ALGO, "Plan file does not exist: ", planPath);
       continue;
     }
     std::string planId = generateFileId(planPath);
@@ -205,9 +227,9 @@ std::shared_ptr<AlgorithmCollection> AlgorithmCollectionBuilder::buildCollection
       auto algoBuilder = std::make_shared<mscclpp::DslAlgorithm>(
           planId, executionPlan, config.tags, mscclpp::Algorithm::Constraint{config.worldSize, config.nRanksPerNode});
       collectionBuilder->addAlgorithmBuilder(algoBuilder);
-      INFO(EXEC, "Successfully loaded plan: ", planId, " for collective: ", config.collective);
+      INFO(ALGO, "Successfully loaded plan: ", planId, " for collective: ", config.collective);
     } catch (const std::exception& e) {
-      WARN(EXEC, "Failed to load plan : ", planPath, " ", e.what());
+      WARN(ALGO, "Failed to load plan : ", planPath, " ", e.what());
     }
   }
   return collection;
@@ -290,7 +312,7 @@ CommResult DslAlgorithm::execute(std::shared_ptr<Communicator> comm, const void*
       executor->execute(rank, (int*)input, (int*)output, inputSize, outputSize, DataType::UINT32, plan_, stream);
       break;
     default:
-      WARN(EXEC, "Unsupported data type: ", static_cast<int>(dtype), " in DslAlgorithm");
+      WARN(ALGO, "Unsupported data type: ", static_cast<int>(dtype), " in DslAlgorithm");
       return CommResult::commInvalidArgument;
   }
   return CommResult::commSuccess;

src/include/logger.hpp

@@ -17,7 +17,7 @@
 namespace mscclpp {
 
 typedef enum : unsigned int { NONE = 0, DEBUG, INFO, WARN, ERROR } LogLevel;
-typedef enum : std::size_t { ENV = 0, NET, CONN, EXEC, NCCL, COUNT } LogSubsys;
+typedef enum : std::size_t { ENV = 0, NET, CONN, EXEC, NCCL, ALGO, COUNT } LogSubsys;
 
 namespace detail {
 std::string guessRemoveProjectPrefix(const std::string& filePathStr);