@JackCaoG and I both have run into cases where our existing concurrency utilities add significant overhead (e.g. waiting >1ms for lock when completing a MultiWait task), which functionally limits the number of threads we can spawn. This PR replaces two custom implementations of common utilities (thread pool and latch/MultiWait) with more optimized upstream equivalents.

• Use tsl::thread::ThreadPool. The underlying implementation more carefully reuses threads and handles NUMA affinity to reduce context switching costs.
  • Simplified interface for starting threads. There used to be two thread pools, one for low-latency ops and one for IO ops. However, there was no difference in implementation between them, so I replaced it with one thread pool with the maximum number of threads. I expect the IO thread pool was more important when we had extremely high latency calls through XRT.
  • A nice side benefit is that by reusing threads more often, our profiles become much more readable:

• Mostly replace MultiWait with absl::BlockingCounter. Completing a task requires acquiring a lock, which can cost >1ms in practice with enough threads. BlockingCounter instead uses a lockless atomic counter, significantly reducing the latency to decrement the remaining task count.
  • MultiWait still exists upstream PyTorch if we need to use it for anything. I left usage of the upstream MultiWait alone.
  • Remove cases of MultiWait where we launch one operation and do not wait for it

We already depend on absl and TSL through OpenXLA, so this PR adds no new dependencies.

Tested with SPMD llama inference to confirm no regression in performance:

Baseline:

Totally decoded 1007 tokens in 7.11816 seconds

With this PR:

Totally decoded 1007 tokens in 7.06291 seconds

I found these performance benefits while working on #5737, which spawns significantly more threads and potentially reduces the synchronous time spent in ExecuteReplicated. See my comment about performance. But, I don't expect this PR to provide significant benefits on its own. Benefits of this PR alone:

• Clearer profile traces, by grouping operations in fewer threads and actually labeling which ones are PyTorch/XLA.
• Deleting code

Future work:

• Check if TSL's low_latency_hint makes a difference in performance. Operations through the thread pool are assumed to be low latency by default.
• See if upstream TSL can remove the tf_ prefix from thread names.
• Use different labeled thread pools for different purposes to make profiles more readable. Example of separating "framework" ops from "runtime" ops in a WIP PR: