mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
[xpu] Support high stream for ProcessGroupXCCL (#163049)
Add high priority stream support for ProcessGroupXCCL. Just like CUDA, XPU streams also support execution with higher priority compared to other streams. Implementation in https://github.com/intel/torch-xpu-ops/pull/1715, add register here. Pull Request resolved: https://github.com/pytorch/pytorch/pull/163049 Approved by: https://github.com/guangyey, https://github.com/gujinghui, https://github.com/EikanWang, https://github.com/albanD
This commit is contained in:
parent
35153d0846
commit
a1005427bf
|
|
@ -3532,11 +3532,35 @@ Example::
|
|||
py::arg("rank"),
|
||||
py::arg("size"),
|
||||
py::arg("options"),
|
||||
R"(Create a new ProcessGroupXCCL instance.)");
|
||||
R"(Create a new ProcessGroupXCCL instance.)")
|
||||
.def(
|
||||
py::init([](const c10::intrusive_ptr<::c10d::Store>& store,
|
||||
int rank,
|
||||
int size) {
|
||||
// gil_scoped_release is not safe as a call_guard in init.
|
||||
// https://github.com/pybind/pybind11/issues/5473
|
||||
py::gil_scoped_release nogil{};
|
||||
|
||||
auto options = ::c10d::ProcessGroupXCCL::Options::create();
|
||||
options->is_high_priority_stream = false;
|
||||
return c10::make_intrusive<::c10d::ProcessGroupXCCL>(
|
||||
store, rank, size, options);
|
||||
}),
|
||||
py::arg("store"),
|
||||
py::arg("rank"),
|
||||
py::arg("size"),
|
||||
R"(Create a new ProcessGroupXCCL instance.)")
|
||||
.def_property_readonly(
|
||||
"options",
|
||||
&::c10d::ProcessGroupXCCL::getOptions,
|
||||
R"(Return the options used to create this ProcessGroupXCCL instance.)");
|
||||
|
||||
intrusive_ptr_class_<::c10d::ProcessGroupXCCL::Options>(
|
||||
processGroupXCCL, "Options", backendOptions)
|
||||
.def(py::init<>());
|
||||
.def(py::init<bool>(), py::arg("is_high_priority_stream") = false)
|
||||
.def_readwrite(
|
||||
"is_high_priority_stream",
|
||||
&::c10d::ProcessGroupXCCL::Options::is_high_priority_stream);
|
||||
module
|
||||
.def(
|
||||
"_dump_xccl_trace",
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user