PR #28877: [XLA]Clamp num_workers to avoid partition overflow #97046
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
6ac02c9 to
27bcefc
Compare
Imported from GitHub PR openxla/xla#28877 When invoking Worker::Parallelize in parallel_for(), the number of workers was previously set to thread_pool_->NumThreads(), which could exceed the number of tasks (n). This mismatch could lead to out-of-bounds access when workers attempt to access partitions that don't exist. This ensures that: - The number of worker threads does not exceed the number of tasks - Each worker maps safely to a valid partition index - Partition and worker ranges stay aligned, preventing undefined behavior Change-Id: I5c3ab319f73009452d27b4db43f7905c5e136c2b Copybara import of the project: -- d72f403a8312b995d5db6112d0e5cb8ad9befff1 by Sharif Inamdar <sharif.inamdar@arm.com>: [XLA]Clamp num_workers to avoid partition overflow When invoking Worker::Parallelize in parallel_for(), the number of workers was previously set to thread_pool_->NumThreads() in oneDNN, which could exceed the number of tasks (n). This mismatch could lead to out-of-bounds access when workers attempt to access partitions that don't exist. Fixed by moving the clamping logic to Worker::Parallelize for all paths. This ensures that: - The number of worker threads does not exceed the number of tasks - Each worker maps safely to a valid partition index - Partition and worker ranges stay aligned, preventing undefined behavior Added simple unit test to test for all usecases. Change-Id: I5c3ab319f73009452d27b4db43f7905c5e136c2b Merging this change closes #28877 PiperOrigin-RevId: 783796500
27bcefc to
fcef98c
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
PR #28877: [XLA]Clamp num_workers to avoid partition overflow
Imported from GitHub PR openxla/xla#28877
When invoking Worker::Parallelize in parallel_for(), the number of workers was previously set to thread_pool_->NumThreads(), which could exceed the number of tasks (n). This mismatch could lead to out-of-bounds access when workers attempt to access partitions that don't exist.
This ensures that:
Change-Id: I5c3ab319f73009452d27b4db43f7905c5e136c2b
Copybara import of the project:
--
d72f403a8312b995d5db6112d0e5cb8ad9befff1 by Sharif Inamdar sharif.inamdar@arm.com:
[XLA]Clamp num_workers to avoid partition overflow
When invoking Worker::Parallelize in parallel_for(), the number
of workers was previously set to thread_pool_->NumThreads() in
oneDNN, which could exceed the number of tasks (n). This mismatch
could lead to out-of-bounds access when workers attempt to access
partitions that don't exist.
Fixed by moving the clamping logic to Worker::Parallelize for
all paths.
This ensures that:
behavior
Added simple unit test to test for all usecases.
Change-Id: I5c3ab319f73009452d27b4db43f7905c5e136c2b
Merging this change closes #28877
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#28877 from isharif168:fix-worker-partition-safety d72f403a8312b995d5db6112d0e5cb8ad9befff1