Use bfdot instruction in ARM bfloat16 dot product if compiling with it available #127488

swolchok · 2024-05-30T00:11:41Z

Stack from ghstack (oldest at bottom):

Improve readability for ARM 16-bit float configuration with to_shift #127971
-> Use bfdot instruction in ARM bfloat16 dot product if compiling with it available #127488

Summary: We should use this intrinsic if the hardware supports it.

Test Plan: Currently untested because I don't have access to appropriate hardware.

[ghstack-poisoned]

pytorch-bot · 2024-05-30T00:11:44Z

🔗 Helpful Links

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

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

❌ 2 New Failures, 6 Unrelated Failures

As of commit bf76663 with merge base 36e9b71 ():

NEW FAILURES - The following jobs have failed:

Check mergeability of ghstack PR / ghstack-mergeability-check (gh)
RuntimeError: PR 127488 is out of sync with the corresponding revision 93a8b3483a6704333ad92dd17a7f45a3c7644e6d on branch gh/swolchok/638/orig that would be merged into main. This usually happens because there is a non ghstack change in the PR. Please sync them and try again (ex. make the changes on origin/gh/swolchok/638/orig and run ghstack).
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
inductor/test_efficient_conv_bn_eval.py::EfficientConvBNEvalCudaTests::test_basic_cuda

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

linux-binary-manywheel / manywheel-py3_8-cuda11_8-test / test (gh) (similar failure)
ImportError: libcudnn.so.8: cannot open shared object file: No such file or directory
trunk / macos-py3-arm64 / test (default, 2, 3, macos-m1-stable) (gh) (similar failure)
'test/dynamo/test_dynamic_shapes.py::DynamicShapesExportTests::test_retracibility_nested_list_out_dynamic_shapes'

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

pull / linux-docs / build-docs-python-false (gh) (trunk failure)
Process completed with exit code 1.

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

linux-binary-manywheel / manywheel-py3_8-cuda12_1-test / test (gh) (#127288)
ImportError: libcudnn.so.8: cannot open shared object file: No such file or directory
linux-binary-manywheel / manywheel-py3_8-cuda12_4-test / test (gh) (#127289)
ImportError: libcudnn.so.8: cannot open shared object file: No such file or directory
pull / linux-focal-cuda12.4-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) ()
inductor/test_efficient_conv_bn_eval.py::EfficientConvBNEvalCudaTests::test_basic_cuda

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

…t available Summary: We should use this intrinsic if the hardware supports it. Test Plan: Currently untested because I don't have access to appropriate hardware. ghstack-source-id: e8b7538 Pull Request resolved: #127488

Skylion007 · 2024-06-01T13:39:51Z

aten/src/ATen/native/BlasKernel.cpp

+#ifdef __ARM_FEATURE_BF16
+  return vbfdotq_f32(a, b, c);
+#else
+  TORCH_CHECK(false, "unsupported operation used!");


Should this be a static_assert False situation? I guess we want the function as a stub even if it's unavailable which precludes that.

Should this be a static_assert False situation?

It should be statically dead in this situation; the TORCH_CHECK is just a backstop in case that isn't true in the future by accident.

[ghstack-poisoned]

…t available Summary: We should use this intrinsic if the hardware supports it. Test Plan: Currently untested because I don't have access to appropriate hardware. ghstack-source-id: 2444407 Pull Request resolved: #127488

[ghstack-poisoned]

aten/src/ATen/native/BlasKernel.cpp

malfet · 2024-06-05T16:23:52Z

aten/src/ATen/native/BlasKernel.cpp

 }

+static C10_ALWAYS_INLINE float32x4_t f32_dot_bf16(float32x4_t a, bfloat16x8_t b, bfloat16x8_t c) {
+#ifdef __ARM_FEATURE_BF16


To trigger it to true one needs to compile with --march=armv8+bf16

aten/src/ATen/native/BlasKernel.cpp

malfet

Hmm, I don't see much perf changes:
Before

mv_nt    torch.float32    4.99 usec
mv_nt    torch.float16    9.70 usec
mv_nt   torch.bfloat16   15.38 usec
mv_ta    torch.float32    3.96 usec
mv_ta    torch.float16   26.92 usec
mv_ta   torch.bfloat16  156.50 usec
notrans  torch.float32    3.96 usec
notrans  torch.float16   28.00 usec
notrans torch.bfloat16   75.61 usec
trans_a  torch.float32    3.94 usec
trans_a  torch.float16   83.89 usec
trans_a torch.bfloat16  366.66 usec
trans_b  torch.float32    7.44 usec
trans_b  torch.float16    9.58 usec
trans_b torch.bfloat16   17.40 usec

after

mv_nt    torch.float32    4.93 usec
mv_nt    torch.float16    9.69 usec
mv_nt   torch.bfloat16   15.18 usec
mv_ta    torch.float32    3.84 usec
mv_ta    torch.float16   26.41 usec
mv_ta   torch.bfloat16  153.94 usec
notrans  torch.float32    3.83 usec
notrans  torch.float16   27.57 usec
notrans torch.bfloat16   74.61 usec
trans_a  torch.float32    3.81 usec
trans_a  torch.float16   83.84 usec
trans_a torch.bfloat16  366.30 usec
trans_b  torch.float32    7.35 usec
trans_b  torch.float16    9.50 usec
trans_b torch.bfloat16   17.38 usec

swolchok · 2024-06-05T17:04:01Z

Hmm, I don't see much perf changes:

very surprising. I would recommend 2 things:

profile with instruments and make sure bfdot is actually getting emitted -- xctrace record --template 'Time Profiler' --launch -- /Users/swolchok/homebrew/Caskroom/miniconda/base/envs/pytorch/bin/python3 benchmarks/benchmark_torch_mm.py llm (replace with path to your python3)
try the larger benchmark cases I sent a PR for: Add LLM benchmark case to benchmark_torch_mm.py malfet/llm_experiments#8

github-actions · 2024-08-04T22:33:59Z

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.

swolchok · 2024-08-26T18:08:13Z

@malfet it looks like your benchmark results above may have been missing 6f18175 (#127318) . Note that your trans_a results match the "Before" and not the "After" on that commit.

swolchok · 2024-08-27T19:51:11Z

getting to the bottom of this now by benchmarking on phone

Port of pytorch/pytorch#127488 . Differential Revision: [D62159105](https://our.internmc.facebook.com/intern/diff/D62159105/) [ghstack-poisoned]

…M_FEATURE_BF16" Port of pytorch/pytorch#127488 . Differential Revision: [D62159105](https://our.internmc.facebook.com/intern/diff/D62159105/) [ghstack-poisoned]

Port of pytorch/pytorch#127488 . Differential Revision: [D62159105](https://our.internmc.facebook.com/intern/diff/D62159105/) [ghstack-poisoned]

…M_FEATURE_BF16" Port of pytorch/pytorch#127488 . Differential Revision: [D62159105](https://our.internmc.facebook.com/intern/diff/D62159105/) [ghstack-poisoned]

Port of pytorch/pytorch#127488 . Differential Revision: [D62159105](https://our.internmc.facebook.com/intern/diff/D62159105/) [ghstack-poisoned]

…M_FEATURE_BF16" Port of pytorch/pytorch#127488 . Differential Revision: [D62159105](https://our.internmc.facebook.com/intern/diff/D62159105/) [ghstack-poisoned]

Port of pytorch/pytorch#127488 . Differential Revision: [D62159105](https://our.internmc.facebook.com/intern/diff/D62159105/) [ghstack-poisoned]

Summary: Pull Request resolved: #5249 Port of pytorch/pytorch#127488 . ghstack-source-id: 242493751 exported-using-ghexport Reviewed By: kimishpatel Differential Revision: D62159105 fbshipit-source-id: e8256ef21373ee5bbc58c15fa2dcf2753af839e0

swolchok · 2024-09-19T18:20:34Z

this will be replaced by a PR to sync the latest work from executorch (including pytorch/executorch#5444 ), and hopefully follow-up to extract a file that can be identical between pytorch/pytorch and pytorch/executorch.

ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes not-yet-reviewed pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) [ghstack-poisoned]

ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes not-yet-reviewed pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) ghstack-source-id: 243596406 Pull Request resolved: #136331

swolchok · 2024-09-19T19:57:28Z

replacement PR is #136331

…t back to ATen BlasKernel" ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes not-yet-reviewed pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) [ghstack-poisoned]

…lasKernel" ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes not-yet-reviewed pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) [ghstack-poisoned]

Pull Request resolved: #136331 ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes not-yet-reviewed pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) ghstack-source-id: 244185360

…t back to ATen BlasKernel" ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes not-yet-reviewed pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) [ghstack-poisoned]

…lasKernel" ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes not-yet-reviewed pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) [ghstack-poisoned]

Pull Request resolved: #136331 ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes not-yet-reviewed pytorch/executorch#5444 . ghstack-source-id: 244187352 Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/)

Pull Request resolved: #136331 ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes pytorch/executorch#5444 . ghstack-source-id: 245455630 Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/)

…t back to ATen BlasKernel" ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) [ghstack-poisoned]

…lasKernel" ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) [ghstack-poisoned]

Pull Request resolved: #136331 ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes pytorch/executorch#5444 . ghstack-source-id: 245903000 Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/)

…t back to ATen BlasKernel" ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) [ghstack-poisoned]

…lasKernel" ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) [ghstack-poisoned]

…136331) ExecuTorch's fork of BlasKernel.cpp grew bfdot support, complete with demonstration that it helps. Port it back to PyTorch. Supersedes #127488 . Includes pytorch/executorch#5444 . Differential Revision: [D63045939](https://our.internmc.facebook.com/intern/diff/D63045939/) Pull Request resolved: #136331 Approved by: https://github.com/malfet, https://github.com/albanD ghstack dependencies: #136445

Pull Request resolved: pytorch/executorch#5249 Port of pytorch/pytorch#127488 . ghstack-source-id: 242493751 @exported-using-ghexport Differential Revision: [D62159105](https://our.internmc.facebook.com/intern/diff/D62159105/)

Update

8f9e48d

[ghstack-poisoned]

swolchok mentioned this pull request May 29, 2024

Extract dot-product functions from fp16_gemv_trans gemv kernels #127435

Closed

swolchok requested a review from malfet May 30, 2024 00:12

swolchok added the topic: performance topic category label May 30, 2024

Skylion007 reviewed Jun 1, 2024

View reviewed changes

Skylion007 approved these changes Jun 1, 2024

View reviewed changes

Update

6b43167

[ghstack-poisoned]

malfet added ciflow/trunk Trigger trunk jobs on your pull request ciflow/linux-aarch64 linux aarch64 CI workflow ciflow/mps Run MPS tests (subset of trunk) labels Jun 3, 2024

Update

a43c328

[ghstack-poisoned]

swolchok mentioned this pull request Jun 4, 2024

Improve readability for ARM 16-bit float configuration with to_shift #127971

Closed

kimishpatel reviewed Jun 4, 2024

View reviewed changes

aten/src/ATen/native/BlasKernel.cpp Show resolved Hide resolved

malfet reviewed Jun 5, 2024

View reviewed changes

aten/src/ATen/native/BlasKernel.cpp Outdated Show resolved Hide resolved

Update aten/src/ATen/native/BlasKernel.cpp

bf76663

malfet reviewed Jun 5, 2024

View reviewed changes

github-actions bot added the Stale label Aug 4, 2024

swolchok mentioned this pull request Sep 10, 2024

[ExecuTorch] Use bfdot if compiled with ARM_FEATURE_BF16 pytorch/executorch#5249

Closed

swolchok closed this Sep 19, 2024

swolchok mentioned this pull request Sep 19, 2024

[PyTorch] Port ExecuTorch bfdot improvement back to ATen BlasKernel #136331

Closed

github-actions bot deleted the gh/swolchok/638/head branch October 20, 2024 02:09

Use bfdot instruction in ARM bfloat16 dot product if compiling with it available #127488

Use bfdot instruction in ARM bfloat16 dot product if compiling with it available #127488

Uh oh!

Conversation

swolchok commented May 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented May 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 2 New Failures, 6 Unrelated Failures

Uh oh!

Skylion007 Jun 1, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

swolchok Jun 3, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

malfet Jun 5, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

malfet left a comment

Choose a reason for hiding this comment

Uh oh!

swolchok commented Jun 5, 2024

Uh oh!

github-actions bot commented Aug 4, 2024

Uh oh!

swolchok commented Aug 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

swolchok commented Aug 27, 2024

Uh oh!

swolchok commented Sep 19, 2024

Uh oh!

swolchok commented Sep 19, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

swolchok commented May 30, 2024 •

edited

Loading

pytorch-bot bot commented May 30, 2024 •

edited

Loading

Skylion007 Jun 1, 2024 •

edited

Loading

malfet Jun 5, 2024 •

edited

Loading

swolchok commented Aug 26, 2024 •

edited

Loading