[xla:cpu] Worker/WorkQueue micro-optimizations

PiperOrigin-RevId: 725314843

Author: ezhulenev

third_party/xla/xla/backends/cpu/runtime/work_queue.h

@@ -31,12 +31,9 @@ limitations under the License.
 #include "absl/container/fixed_array.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
-#include "absl/time/time.h"
 #include "xla/tsl/concurrency/async_value_ref.h"
 #include "xla/tsl/concurrency/chain.h"
 #include "xla/tsl/lib/math/math_util.h"
-#include "xla/tsl/platform/env.h"
-#include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/logging.h"
 
 #define EIGEN_USE_THREADS
@@ -80,7 +77,7 @@ class WorkQueue {
   };
 
   absl::FixedArray<Partition, 32> partitions_;
-  alignas(kAtomicAlignment) std::atomic<size_t> empty_;
+  alignas(kAtomicAlignment) std::atomic<bool> empty_;
 };
 
 // Worker processes tasks from the work queue starting from the assigned
@@ -165,10 +162,7 @@ inline std::optional<size_t> Worker::Pop() {
   std::optional<size_t> task = queue_->Pop(partition_index_);
   if (task) return task;
 
-  // If work queue is empty, we are not going to find any more tasks.
-  if (queue_->empty()) return std::nullopt;
-
-  while (!task.has_value()) {
+  while (!task.has_value() && !queue_->empty()) {
     // Wrap around to the first partition.
     if (ABSL_PREDICT_FALSE(++partition_index_ >= queue_->num_partitions())) {
       partition_index_ = 0;