[kernels] Migrate comm/ from LegacyUnsafePointer to UnsafePointer#5690
[kernels] Migrate comm/ from LegacyUnsafePointer to UnsafePointer#5690KrxGu wants to merge 4 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR migrates the comm/ directory from the legacy LegacyUnsafePointer API to the new UnsafePointer API with explicit mutability and origin parameters. The changes focus on three core communication kernel files for multi-GPU operations.
Key Changes:
- Replaced
LegacyUnsafePointerwithUnsafePointerand added explicitmut=True/Falseparameters throughout - Replaced
LegacyOpaquePointerwithMutOpaquePointer[MutOrigin.external]for FFI callback boundaries - Removed the temporary
UnsafePointerV2workaround and consolidated to the newUnsafePointerAPI
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| max/kernels/src/comm/sync.mojo | Updated imports, migrated OpaquePointer types for FFI callbacks, added explicit mutability to Signal pointer types in barrier synchronization functions |
| max/kernels/src/comm/allreduce.mojo | Migrated pointer types across all reduction kernels and helper functions, added explicit mutability for source (mut=False) and destination (mut=True) buffers, updated signal pointer types |
| max/kernels/src/comm/allgather.mojo | Updated pointer types in allgather kernels, added explicit mutability for output (mut=True) and source (mut=False) buffers, migrated signal pointer types |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
@NathanSWard This PR(Part-1) is ready for review. Ci all green, Good To Merge |
NathanSWard
left a comment
NathanSWard
left a comment
There was a problem hiding this comment.
Thanks for tackling this!
Let a handful of comments - mostly around when needing/not needing to manually specify an origin or mut parameter.
Hopefully it makes sense, but if not please ask questions and I'll do my best to help/explain!
|
@NathanSWard i have addressed all your reviews and made the necessary changes. CI is all green, Good To Merge!! |
|
@KrxGu - thanks! I'll take a look in the next few days! Modular is on a Holiday break but I'll get this is soon :) |
Nw i can wait😄,Thank you for getting back!! |
|
@KrxGu - now that we're back from holidays I'm re-prioritizing getting your PRs merged in! |
NathanSWard
left a comment
NathanSWard
left a comment
There was a problem hiding this comment.
After this gets rebased and the last to comments are address this looks great and I'll merge it in :D
Replace LegacyUnsafePointer with the new UnsafePointer API in sync.mojo, allgather.mojo, and allreduce.mojo. Key changes: - Use explicit mut=True/False and origin parameters - Update OpaquePointer to MutOpaquePointer for FFI callbacks - MutAnyOrigin for pointer types that need to work across origins Part of modular#5673
- Add explicit mut=True/False to pointer signatures in allreduce.mojo - Migrate distributed_matmul.mojo from LegacyUnsafePointer to UnsafePointer - Update all rank_sigs type signatures to use new pointer API
Use RealUnsafePointer alias for rank_sigs in allreduce, allgather, and matmul_allreduce calls while keeping LegacyUnsafePointer for existing usages in the file.
- Remove unnecessary imports (types are in prelude) - Use ImmutAnyOrigin for immutable pointers (mut=False) - Remove mut= parameter when origin is explicit - Keep origins explicit for GPU kernel parameters (enqueue_function_checked) - Simplify pointer type signatures across all files
KrxGu
force-pushed
the
krxgu/migrate-comm-shmem-unsafe-pointer-part1
branch
from
cc579c4 to
70a7bae
Compare
|
The CI failures in @NathanSWard Kindly confirm that the above is the case, if not i am ready to make any changes suggested by you. |
|
Correct @KrxGu - these issues are orthogonal to your changes. You can ignore them! |
Gotcha, so now this PR is Good To Merge!! |
|
!sync |
modularbot
added
the
imported-internally
|
✅🟣 This contribution has been merged 🟣✅ Your pull request has been merged to the internal upstream Mojo sources. It will be reflected here in the Mojo repository on the main branch during the next Mojo nightly release, typically within the next 24-48 hours. We use Copybara to merge external contributions, click here to learn more. |
modularbot
added
merged-internally
|
Landed in 4f80969! Thank you for your contribution 🎉 |
Fixes #5673
Summary
Migrates the comm module and its dependents from
LegacyUnsafePointerto the newUnsafePointerAPI with explicit mutability and origin parameters.Part 1 of #5673.
Changes
comm/ module:
sync.mojo: Updated imports, replacedLegacyOpaquePointerwithMutOpaquePointer[MutOrigin.external]for FFI callbacksallgather.mojo: Added explicitmutandMutAnyOriginto all pointer signaturesallreduce.mojo: Migrated all pointer usages to new APIDependent files:
linalg/distributed_matmul.mojo: Updated import andrank_sigssignaturesMogg/MOGGKernelAPI/MOGGKernelAPI.mojo: AddedRealUnsafePointeralias for comm API calls while preservingLegacyUnsafePointerfor other usagesTesting
Builds successfully:
./bazelw build //max/kernels/src/comm:comm./bazelw build //max/kernels/src/linalg:linalg./bazelw build //max/kernels/src/Mogg/MOGGKernelAPI:MOGGKernelAPINext Steps
Follow-up PRs will migrate
shmem/andcomm/vendor/ccl.mojo.