Parallelize optimized op_log_softmax

[swolchok]

Stack from ghstack (oldest at bottom):

• -> Parallelize optimized op_log_softmax #12099
• Use shared log_softmax kernels from PyTorch #12098

Straightforward application of parallel_for.

Differential Revision: D76831122

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/12099

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 2 Cancelled Jobs

As of commit bd32947 with merge base bed504e ():

NEW FAILURE - The following job has failed:

• Propose to merge ghstack orig PRs to main / Try to create a PR with ghstack /orig branch (gh)
  Process completed with exit code 1.

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• pull / test-models-linux (mobilebert, portable, linux.2xlarge) / linux-job (gh)
  ##[error]The operation was canceled.
• pull / test-models-linux (mobilebert, xnnpack-quantization-delegation, linux.2xlarge) / linux-job (gh)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D76831122

[swolchok]

pull / android / run-emulator is consistently telling us this is broken, but the bug is not obvious to me on rereading

[swolchok]

after rereading the code twice, I was inclined to suspect that the kernel is somehow thread-unsafe, so I built with TSAN locally (-fsanitize=debug, CMAKE_BUILD_TYPE=RelWithDebInfo so parallel_for will actually parallelize) and didn't catch a failure. doesn't look like a pre-existing issue on HUD or anything so we shouldn't bypass, but I'm not sure what to do next. Nobody is particularly asking for parallelization AFAIK so there's no rush...

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D76831122

[swolchok]

android / run-emulator is no longer complaining. I guess that means this is landable!

[swolchok]

mobilebert has been timing out on HUD, bypassing