Remove Stream::SetPriority functions from Stream.

They don't work after the stream is initialized in GpuStream (the only Stream implementation to make use of the priority).  Instead, move the parameter to Stream::Initialize, which is the only place it's actually used.

PiperOrigin-RevId: 622958008

third_party/xla/xla/stream_executor/stream.cc

@@ -19,6 +19,7 @@ limitations under the License.
 #include <cstdint>
 #include <functional>
 #include <memory>
+#include <optional>
 #include <sstream>
 #include <string>
 #include <utility>
@@ -59,14 +60,6 @@ Stream::~Stream() {
   }
 }
 
-void Stream::SetPriority(StreamPriority priority) {
-  implementation_->SetPriority(priority);
-}
-
-void Stream::SetPriority(int priority) {
-  implementation_->SetPriority(priority);
-}
-
 std::variant<StreamPriority, int> Stream::priority() const {
   return implementation_->priority();
 }
@@ -88,7 +81,8 @@ absl::Status Stream::RefreshStatus() {
   return status;
 }
 
-absl::Status Stream::Initialize() {
+absl::Status Stream::Initialize(
+    std::optional<std::variant<StreamPriority, int>> priority) {
   absl::MutexLock lock(&mu_);
   if (allocated_) {
     return absl::InternalError(
@@ -98,7 +92,13 @@ absl::Status Stream::Initialize() {
     return absl::InternalError(
         "stream should be in !ok() state pre-initialization");
   }
-
+  if (priority.has_value()) {
+    if (std::holds_alternative<StreamPriority>(*priority)) {
+      implementation_->SetPriority(std::get<StreamPriority>(*priority));
+    } else {
+      implementation_->SetPriority(std::get<int>(*priority));
+    }
+  }
   if (parent_->AllocateStream(this)) {
     // Successful initialization!
     allocated_ = true;

third_party/xla/xla/stream_executor/stream.h

@@ -23,6 +23,7 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <variant>
@@ -110,7 +111,8 @@ class Stream {
 
   // Initialize the stream. This must be performed before entraining any other
   // operations.
-  absl::Status Initialize();
+  absl::Status Initialize(
+      std::optional<std::variant<StreamPriority, int>> priority = std::nullopt);
 
   // Get or create a sub-stream from this stream. If there is any sub-stream in
   // the pool that can be reused then just return this sub-stream.  Otherwise
@@ -292,9 +294,6 @@ class Stream {
     return parent()->GetDeviceDescription().rocm_compute_capability();
   }
 
-  void SetPriority(StreamPriority priority);
-  void SetPriority(int priority);
-
   std::variant<StreamPriority, int> priority() const;
 
  private:

third_party/xla/xla/stream_executor/stream_executor_pimpl.cc

@@ -442,14 +442,7 @@ Event::Status StreamExecutor::PollForEventStatus(Event* event) {
 absl::StatusOr<std::unique_ptr<Stream>> StreamExecutor::CreateStream(
     std::optional<std::variant<StreamPriority, int>> priority) {
   auto stream = std::make_unique<Stream>(this);
-  if (priority.has_value()) {
-    if (std::holds_alternative<StreamPriority>(*priority)) {
-      stream->SetPriority(std::get<StreamPriority>(*priority));
-    } else {
-      stream->SetPriority(std::get<int>(*priority));
-    }
-  }
-  TF_RETURN_IF_ERROR(stream->Initialize());
+  TF_RETURN_IF_ERROR(stream->Initialize(priority));
   return std::move(stream);
 }