Use dynamic values for AtomicOps #269
Merged
Conversation
This file contains hidden or 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
panditsa
force-pushed
the
sanket/dyn_atomicop
branch
2 times, most recently
from
458218d to
27db7bf
Compare
panditsa
force-pushed
the
sanket/dyn_atomicop
branch
2 times, most recently
from
3f52c02 to
e989252
Compare
panditsa
force-pushed
the
sanket/dyn_atomicop
branch
2 times, most recently
from
ba87142 to
c96ada3
Compare
panditsa
force-pushed
the
sanket/dyn_atomicop
branch
2 times, most recently
from
c66c964 to
66c71f0
Compare
panditsa
force-pushed
the
sanket/dyn_atomicop
branch
4 times, most recently
from
2fd2856 to
09a7262
Compare
willghatch
reviewed
Sep 26, 2025
|
|
||
| elements_per_thread: Optional[Any] = None | ||
| mapping: Optional[IndexMapping] = None | ||
| mapping_dynamic_vals: tuple[fx.Node, ...] = () |
Contributor
There was a problem hiding this comment.
Unrelated to this PR exactly, but why does AtomicOp need to inherit from both BinaryOpBase and ABC, when BinaryOpBase has ABC in its inheritance chain? Just curious.
Contributor
Author
There was a problem hiding this comment.
good question, there are a couple of other classes with the same dependency. But thanks for pointing it out, I removed it for AtomicOp and everything seems fine.
willghatch
approved these changes
Sep 26, 2025
Contributor
willghatch
left a comment
willghatch
left a comment
There was a problem hiding this comment.
This looks reasonable to me.
panditsa
force-pushed
the
sanket/dyn_atomicop
branch
2 times, most recently
from
99d7cd5 to
f6b50a3
Compare
Signed-off-by: Sanket Pandit <sanketp@amd.com> Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Signed-off-by: Sanket Pandit <sanketp@amd.com> Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
saladpalad
pushed a commit
to saladpalad/wave
that referenced
this pull request
Oct 3, 2025
Add dynamic indexing support for atomic operations in MoE alignment For the MoE align block, we need support for dynamically indexed values in atomic operations. This change introduces `mapping_dynamic_vals` parameter to AtomicOp and updates the atomic operation handler to support dynamic indexing. Key changes: - Add `mapping_dynamic_vals` parameter to atomic operations - Update `handle_atomic_op` decorator to process dynamic values - Convert dynamic registers to index-typed vectors and extract scalar indices - Pass dynamic value dictionary to `_build_start_indices` for proper index computation This enables atomic operations like `atomic_add` to use runtime-computed indices from input data --------- Signed-off-by: Sanket Pandit <sanketp@amd.com> Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Megan0704-1
pushed a commit
to Megan0704-1/wave
that referenced
this pull request
Oct 28, 2025
Add dynamic indexing support for atomic operations in MoE alignment For the MoE align block, we need support for dynamically indexed values in atomic operations. This change introduces `mapping_dynamic_vals` parameter to AtomicOp and updates the atomic operation handler to support dynamic indexing. Key changes: - Add `mapping_dynamic_vals` parameter to atomic operations - Update `handle_atomic_op` decorator to process dynamic values - Convert dynamic registers to index-typed vectors and extract scalar indices - Pass dynamic value dictionary to `_build_start_indices` for proper index computation This enables atomic operations like `atomic_add` to use runtime-computed indices from input data --------- Signed-off-by: Sanket Pandit <sanketp@amd.com> Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Add dynamic indexing support for atomic operations in MoE alignment
For the MoE align block, we need support for dynamically indexed values in atomic operations. This change introduces
mapping_dynamic_valsparameter to AtomicOp and updates the atomic operation handler to support dynamic indexing.Key changes:
mapping_dynamic_valsparameter to atomic operationshandle_atomic_opdecorator to process dynamic values_build_start_indicesfor proper index computationThis enables atomic operations like
atomic_addto use runtime-computed indices from input data