[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:
Han Chao 2025-10-22 00:54:23 +00:00 committed by PyTorch MergeBot
parent 35153d0846
commit a1005427bf

View File

@ -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",