[Neuron] Fix the XLADevice Neuron mappings for SPMD downcasts #8335
+23
−8
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
aws-rhsoln
reviewed
Oct 29, 2024
rpsilva-aws
force-pushed
the
neuron_device_mapping
branch
from
418c0c0
to
145d3d5
Compare
jeffhataws
approved these changes
Oct 30, 2024
aws-nm9
approved these changes
Oct 30, 2024
aws-kingrj
reviewed
Oct 30, 2024
| return xla::PrimitiveType::F32; | ||
| } | ||
| return xla::PrimitiveType::F64; | ||
| case xla::PrimitiveType::F32: | ||
| return UseBF16() || DowncastBF16() ? xla::PrimitiveType::BF16 | ||
| : xla::PrimitiveType::F32; | ||
| case xla::PrimitiveType::U16: | ||
| return hw_type != XlaDeviceType::NEURON ? xla::PrimitiveType::U16 | ||
| : xla::PrimitiveType::U32; | ||
| return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::U32 |
There was a problem hiding this comment.
Does the deprecating of the XLA casting env variables affect this at all?
There was a problem hiding this comment.
No, this change only affects the check for the device type, not the downcast flag check(or the deprecation of those flags). Since the pre-existing logic incorporated the deprecated flags, the current one does the same.
The flags being deprecated will not affect this downcasting check change.
aws-kingrj
approved these changes
Oct 30, 2024
| return xla::PrimitiveType::F32; | ||
| } | ||
| return xla::PrimitiveType::F64; | ||
| case xla::PrimitiveType::F32: | ||
| return UseBF16() || DowncastBF16() ? xla::PrimitiveType::BF16 | ||
| : xla::PrimitiveType::F32; | ||
| case xla::PrimitiveType::U16: | ||
| return hw_type != XlaDeviceType::NEURON ? xla::PrimitiveType::U16 | ||
| : xla::PrimitiveType::U32; | ||
| return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::U32 |
qihqi
approved these changes
Oct 31, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Currently, the downcasting mappings based on the XLA device types does not account for virtual devices (SPMD), when requiring Neuron specific casts. The same applies to the lowering and native functions. This PR fixes it by deriving the proper hardware device.