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

As of commit a04da54 with merge base cf0bfd2 ():

NEW FAILURE - The following job has failed:

• pull / android / run-emulator (gh)
  The process '/usr/bin/sh' failed with exit code 255

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...