Optimize GELU BFloat16 Impl in CPU path

[yanbing-j]

Description

For slow path (with non-contiguous inputs) with none or tanh approximate, current bfloat16 impl is not performance friendly in ATen. This PR uses float32 as an immediate type, in order to reduce the heavy cost of converting bf16 to fp32.

Test

IceLake 2S 32C (Intel(R) Xeon(R) Platinum 8358 CPU @ 2.60GHz)

single socket (32 cores):
approximate is none:

input shapes	forward ( base) (ms)	backward (base) (ms)	forward (optimized) (ms)	backward (optimized) (ms)
[16, 32, 32]	0.361	1.055	0.348	0.672
[32, 32, 64]	0.084	2.003	0.076	1.426
[32, 64, 128]	0.237	2.007	0.22	1.454
[64, 128, 128]	2.23	6.348	1.943	4.103

approximate is tanh:

input shapes	forward ( base) (ms)	backward (base) (ms)	forward (optimized) (ms)	backward (optimized) (ms)
[16, 32, 32]	0.203	1.209	0.138	0.474
[32, 32, 64]	0.063	2.497	0.043	0.985
[32, 64, 128]	0.201	2.707	0.141	1.205
[64, 128, 128]	1.549	8.749	1.065	3.635

single core:
approximate is none:

input shapes	forward ( base) (ms)	backward (base) (ms)	forward (optimized) (ms)	backward (optimized) (ms)
[16, 32, 32]	0.359	1.055	0.267	0.592
[32, 32, 64]	1.11	3.483	1.063	2.373
[32, 64, 128]	4.478	13.866	4.27	9.426
[64, 128, 128]	17.675	55.231	16.805	37.509

approximate is tanh:

input shapes	forward ( base) (ms)	backward (base) (ms)	forward (optimized) (ms)	backward (optimized) (ms)
[16, 32, 32]	0.202	1.212	0.138	0.473
[32, 32, 64]	0.776	4.843	0.531	1.872
[32, 64, 128]	3.203	19.267	2.16	7.243
[64, 128, 128]	12.33	76.834	8.286	29.553

cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10

[facebook-github-bot]

🔗 Helpful links

• 🧪  See artifacts and rendered test results at hud.pytorch.org/pr/79378
• 📄  Preview Python docs built from this PR
• 📄  Preview C++ docs built from this PR
• ❓Need help or want to give feedback on the CI? Visit our office hours

✅ No Failures (0 Pending)

As of commit a917806 (more details on the Dr. CI page):

Expand to see more

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

[facebook-github-bot]

/easycla

As part of the transition to the PyTorch Foundation, this project now requires contributions be covered under the new CLA. See #85559 for additional details.

This comment will trigger a new check of this PR. If you are already covered, you will simply see a new "EasyCLA" check that passes. If you are not covered, a bot will leave a new comment with a link to sign.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: yanbing-j / name: YanbingJiang (e1e40a0, 6df4c71, d1dd927, edb2202, 19fc33a, d28905e, a917806, 84e1a5d)

[yanbing-j]

@pytorchbot label ciflow/trunk

[pytorch-bot]

Can't add following labels to PR: ciflow/trunk Please ping one of the reviewers for help.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

✅ No Failures

As of commit ddda96a:
💚 Looks good so far! There are no failures yet. 💚

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

[mingfeima]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here