Redo add parallel_for to torch/csrc/stable #166695
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/166695
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 1 New FailureAs of commit b8cb398 with merge base 4295a9a ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
| const int64_t end, | ||
| const int64_t grain_size, | ||
| const F& f) { | ||
| if (end - begin < grain_size) { |
There was a problem hiding this comment.
@malfet want to double check this matches what you asked for on chat (I interpreted that you meant end - start < grain_size)
May be the logic we want to preserve in header only if range is trivial (i.e. end-start < 1) we just run this thing sequentially
If I misunderstood and you actually meant we want to preserve this in the header I'll redo this
pytorch/aten/src/ATen/Parallel-inl.h
Lines 23 to 31 in 11f73d7
There was a problem hiding this comment.
IMO this is an optimization that could be landed later, and may be not needed here, i.e. often when dev calls at::parallel_for their regions should be reasonably large
| const int64_t grain_size, | ||
| const F& f) { | ||
| if (end - begin < grain_size) { | ||
| f(begin, end); |
There was a problem hiding this comment.
Also I didn't add this as I wasn't sure if it's actually needed
internal::ThreadIdGuard tid_guard(0);
c10::ParallelGuard guard(true);
Let me know if I should shim and add these here
| int64_t sizes[] = {size}; | ||
| int64_t strides[] = {1}; |
There was a problem hiding this comment.
you might find that you get internal lint complaints about C arrays because of this. Since you specifically need an array of size 1, you can just not make these arrays:
int64_t stride = 1;
aoti_torch_empty_strided(
1,
&size,
&stride,
// ...
|
|
||
| // Get the current thread index in a parallel region | ||
| // Returns 0 if not in a parallel region | ||
| AOTI_TORCH_EXPORT int32_t torch_get_thread_idx(); |
There was a problem hiding this comment.
Is there a document on signed vs unsigned types? I.e. what thread_idx == -1 means?
Is there an expectations that it should be in range [0, torch_get_max_threads())?
| const int64_t end, | ||
| const int64_t grain_size, | ||
| const F& f) { | ||
| if (end - begin < grain_size) { |
There was a problem hiding this comment.
IMO this is an optimization that could be landed later, and may be not needed here, i.e. often when dev calls at::parallel_for their regions should be reasonably large
Stack from ghstack (oldest at bottom):