Skip to content

[Compile] Add NEON implementation of fp32->bf16 conversion #137131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
malfet wants to merge 1 commit into from
Closed

[Compile] Add NEON implementation of fp32->bf16 conversion #137131

malfet wants to merge 1 commit into from

Conversation

@malfet
Copy link
Contributor

@malfet malfet commented Oct 1, 2024 •
edited by pytorch-bot bot
Loading

@malfet
Update
b6ef72d 
[ghstack-poisoned]
@malfet malfet mentioned this pull request Oct 1, 2024
Closed
@pytorch-bot
Copy link

pytorch-bot bot commented Oct 1, 2024 •
edited
Loading

🔗 Helpful Links

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

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

✅ No Failures

As of commit b6ef72d with merge base b1b6816 (image):
💚 Looks good so far! There are no failures yet. 💚

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

@pytorch-bot pytorch-bot bot added the module: cpu CPU specific problem (e.g., perf, algorithm) label Oct 1, 2024
malfet added a commit that referenced this pull request Oct 1, 2024
@malfet
[Compile] Add NEON implementation of fp32->bf16 conversion
522a5f3 
ghstack-source-id: 10b7e8c
Pull Request resolved: #137131
@malfet malfet requested review from jgong5 and kimishpatel October 2, 2024 14:22
kimishpatel
kimishpatel requested changes Oct 2, 2024
View reviewed changes
Copy link
Contributor

@kimishpatel kimishpatel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

left some comments to emit fewer instructions

aten/src/ATen/cpu/vec/vec256/vec256_convert.h
VectorizedN<BFloat16, 1> result;
uint32x4x2_t u32x4x2_0 = vld1q_u32_x2(reinterpret_cast<const uint32_t*>(&src[0]));
uint32x4x2_t u32x4x2_1 = vld1q_u32_x2(reinterpret_cast<const uint32_t*>(&src[1]));
uint16x4_t u16x4_0 = vmovn_u32(vshrq_n_u32(u32x4x2_0.val[0], 16));
Copy link
Contributor

@kimishpatel kimishpatel Oct 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
uint16x4_t u16x4_0 = vmovn_u32(vshrq_n_u32(u32x4x2_0.val[0], 16));
uint16x4_t u16x4_0 =vshrn_n_u32(u32x4x2_0.val[0], 16);

You can use narrowing shift to avoid vmovn. ARM is extremely rich in types of instructions it supports but it is also confusing because there can always be more compact way of doing that.

aten/src/ATen/cpu/vec/vec256/vec256_convert.h
uint32x4x2_t u32x4x2_0 = vld1q_u32_x2(reinterpret_cast<const uint32_t*>(&src[0]));
uint32x4x2_t u32x4x2_1 = vld1q_u32_x2(reinterpret_cast<const uint32_t*>(&src[1]));
uint16x4_t u16x4_0 = vmovn_u32(vshrq_n_u32(u32x4x2_0.val[0], 16));
uint16x4_t u16x4_1 = vmovn_u32(vshrq_n_u32(u32x4x2_0.val[1], 16));
Copy link
Contributor

@kimishpatel kimishpatel Oct 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
uint16x4_t u16x4_1 = vmovn_u32(vshrq_n_u32(u32x4x2_0.val[1], 16));
uint16x8_t u16x8_low =vshrn_high_n_u32(u16x4_0, u32x4x2_0.val[1], 16);

and this can be used to avoid vcombine

kimishpatel
kimishpatel requested changes Oct 2, 2024
View reviewed changes
Copy link
Contributor

@kimishpatel kimishpatel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

left some comments to emit fewer instructions

jgong5
jgong5 reviewed Oct 8, 2024
View reviewed changes
aten/src/ATen/cpu/vec/vec256/vec256_convert.h
Comment on lines +319 to +321
struct VecConvert<BFloat16, 1, float, 2> {
static inline VectorizedN<BFloat16, 1> apply(
const VectorizedN<float, 2>& src) {
Copy link
Collaborator

@jgong5 jgong5 Oct 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

will you add "1 bf16" -> "2 fp32" too?

@github-actions
Copy link
Contributor

github-actions bot commented Dec 7, 2024

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.

@github-actions github-actions bot added the Stale label Dec 7, 2024
@github-actions github-actions bot closed this Jan 6, 2025
@github-actions github-actions bot deleted the gh/malfet/31/head branch February 9, 2025 02:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jgong5 jgong5 jgong5 left review comments

@kimishpatel kimishpatel kimishpatel requested changes

Assignees

No one assigned

Labels

ciflow/linux-aarch64 linux aarch64 CI workflow module: cpu CPU specific problem (e.g., perf, algorithm) release notes: inductor Stale topic: performance topic category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@malfet @jgong5 @kimishpatel