mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-04-19 14:29:13 +00:00
Torch integration (#692)
Reorganize current native algorithm implementation and DSL algorithm implementation. Provide unified API for DSL algo and native algo and provide interface to tune the algo Provide interface for pytorch integration with native API and DSL --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: chhwang <8018170+chhwang@users.noreply.github.com>
This commit is contained in:
@@ -15,50 +15,8 @@ namespace nb = nanobind;
|
||||
using namespace mscclpp;
|
||||
|
||||
void register_executor(nb::module_& m) {
|
||||
nb::enum_<DataType>(m, "DataType")
|
||||
.value("int32", DataType::INT32)
|
||||
.value("uint32", DataType::UINT32)
|
||||
.value("float16", DataType::FLOAT16)
|
||||
.value("float32", DataType::FLOAT32)
|
||||
.value("bfloat16", DataType::BFLOAT16);
|
||||
|
||||
nb::enum_<PacketType>(m, "PacketType").value("LL8", PacketType::LL8).value("LL16", PacketType::LL16);
|
||||
|
||||
nb::class_<ExecutionRequest>(m, "ExecutionRequest")
|
||||
.def_ro("world_size", &ExecutionRequest::worldSize)
|
||||
.def_ro("n_ranks_per_node", &ExecutionRequest::nRanksPerNode)
|
||||
.def_prop_ro(
|
||||
"input_buffer",
|
||||
[](const ExecutionRequest& self) -> uintptr_t { return reinterpret_cast<uintptr_t>(self.inputBuffer); })
|
||||
.def_prop_ro(
|
||||
"output_buffer",
|
||||
[](const ExecutionRequest& self) -> uintptr_t { return reinterpret_cast<uintptr_t>(self.outputBuffer); })
|
||||
.def_ro("message_size", &ExecutionRequest::messageSize)
|
||||
.def_prop_ro("collective", [](ExecutionRequest& self) -> const std::string& { return self.collective; })
|
||||
.def_prop_ro("hints", [](ExecutionRequest& self) { return self.hints; });
|
||||
|
||||
nb::class_<ExecutionPlanHandle>(m, "ExecutionPlanHandle")
|
||||
.def_ro("id", &ExecutionPlanHandle::id)
|
||||
.def_ro("constraint", &ExecutionPlanHandle::constraint)
|
||||
.def_ro("plan", &ExecutionPlanHandle::plan)
|
||||
.def_ro("tags", &ExecutionPlanHandle::tags)
|
||||
.def_static("create", &ExecutionPlanHandle::create, nb::arg("id"), nb::arg("world_size"),
|
||||
nb::arg("nranks_per_node"), nb::arg("plan"),
|
||||
nb::arg("tags") = std::unordered_map<std::string, uint64_t>{});
|
||||
|
||||
nb::class_<ExecutionPlanHandle::Constraint>(m, "ExecutionPlanConstraint")
|
||||
.def_ro("world_size", &ExecutionPlanHandle::Constraint::worldSize)
|
||||
.def_ro("n_ranks_per_node", &ExecutionPlanHandle::Constraint::nRanksPerNode);
|
||||
|
||||
nb::class_<ExecutionPlanRegistry>(m, "ExecutionPlanRegistry")
|
||||
.def_static("get_instance", &ExecutionPlanRegistry::getInstance)
|
||||
.def("register_plan", &ExecutionPlanRegistry::registerPlan, nb::arg("planHandle"))
|
||||
.def("get_plans", &ExecutionPlanRegistry::getPlans, nb::arg("collective"))
|
||||
.def("get", &ExecutionPlanRegistry::get, nb::arg("id"))
|
||||
.def("set_selector", &ExecutionPlanRegistry::setSelector, nb::arg("selector"))
|
||||
.def("set_default_selector", &ExecutionPlanRegistry::setDefaultSelector, nb::arg("selector"))
|
||||
.def("clear", &ExecutionPlanRegistry::clear);
|
||||
|
||||
nb::class_<ExecutionPlan>(m, "ExecutionPlan")
|
||||
.def(nb::init<const std::string&, int>(), nb::arg("planPath"), nb::arg("rank"))
|
||||
.def_prop_ro("name", [](const ExecutionPlan& self) -> std::string { return self.name(); })
|
||||
|
||||
Reference in New Issue
Block a user