[XLA:CPU] Possible perf regression in `Worker::Parallelize`

Apologies if I'm misreading the code but I'd appreciate someone with context confirming or telling me I've missed something. ## Observation At [`work_queue.h:240`](https://github.com/openxla/xla/blob/main/xla/backends/cpu/runtime/work_queue.h#L240): ```cpp work_queue(num_work_items, /*num_partitions=*/this->count_down.count()) ``` In the original [9d95a71119](https://github.com/openxla/xla/commit/9d95a71119), `count_down` counted **workers**, so this line correctly meant `num_partitions = num_workers`. Then [5db7c22b15](https://github.com/openxla/xla/commit/5db7c22b15) ("Make Worker::Parallelize deadlock-proof") repurposed `count_down` to count **items**: ```diff - tsl::CountDownAsyncValueRef<tsl::Chain> count_down(num_workers); + tsl::CountDownAsyncValueRef<tsl::Chain> count_down(num_tasks); ``` which makes sense as a correctness fix, but unless I'm missing something the partition-count line wasn't updated. `count_down.count()` now returns `num_tasks`, so we end up with one `ABSL_CACHELINE_ALIGNED` partition per item: 64B each, `FixedArray<Partition, 32>` spills to heap for any job > 32 items, and every `Pop` walks to a fresh cache line via work-stealing wrap-around. Could be deliberate for stealing granularity, but the line reads like it intended the old semantics. ## Impact Added a small `BM_WorkerParallelize` micro (noop body, AMD EPYC 9634 SMT, 3 trials × 20 reps × min-time 0.3 s): | `num_work_items` | `partitions = items` (current) | `partitions = workers` | diff | |---:|---:|---:|---:| | 128 | 45,412 ns | 16,251 ns | −64% | | 256 | 65,045 ns | 22,824 ns | −65% | | 512 | 74,796 ns | 35,227 ns | −53% | | 1024 | 82,599 ns | 52,293 ns | −37% | ## Proposed Fix Substitute `num_workers` within `ParallelizeContext`: ```diff count_down(std::move(count_down)), - work_queue(num_work_items, /*num_partitions=*/this->count_down.count()), + work_queue(num_work_items, /*num_partitions=*/num_workers), parallel_work(std::forward<ParallelWork>(parallel_work)) {} ``` `WorkQueue` already supports `num_partitions != num_work_items` - [`work_queue_test.cc:43-81`](https://github.com/openxla/xla/blob/main/xla/backends/cpu/runtime/work_queue_test.cc#L43) covers `(14 items, 4 partitions)` etc., so I don't believe this breaks anything. Happy to send a PR with the fix and the micro if this is wanted. Filing the issue first in case I've misread the design intent. Perhaps @ezhulenev best positioned / most familiar with change?