[Kernels] Migrate CPU matmul kernels from LegacyUnsafePointer to UnsafePointer #5685
Conversation
Updated 6 CPU matmul kernel files to use the new UnsafePointer API with explicit mutability and origin tracking. The migration uses interop with LegacyUnsafePointer where needed for compatibility with unmigrated code like _Accumulator and FFI interfaces. Fixes modular#5670
There was a problem hiding this comment.
Pull request overview
This PR migrates CPU matmul kernel implementations from the deprecated LegacyUnsafePointer alias to the new UnsafePointer API across 6 files in the matmul CPU kernel directory.
Key changes:
- Updated imports to use
LegacyUnsafePointerexplicitly instead of aliasing it asUnsafePointer - Updated function signatures to use the new
UnsafePointerAPI with explicit mutability and origin tracking - Added
rebindconversions where needed for compatibility with unmigrated dependencies and FFI interfaces
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| vnni.mojo | Updated import statement and rebind operations to convert new UnsafePointer to LegacyUnsafePointer for _Accumulator compatibility |
| neon.mojo | Updated import statement and rebind operations to convert new UnsafePointer to LegacyUnsafePointer for _Accumulator compatibility |
| impl.mojo | Updated imports and function signature for _view_buffer_as to accept new UnsafePointer API; updated variable declarations |
| i8mm.mojo | Added both imports, updated LoadStore_i8mm function signatures to new API, but contains incorrect rebind operations at call sites |
| default.mojo | Updated import statement and rebind operations to convert new UnsafePointer to LegacyUnsafePointer for _Accumulator compatibility |
| apple_accelerate.mojo | Removed UnsafePointer alias import, updated type definitions and rebind operations for FFI compatibility with CBLAS functions |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
The c_ptr from LayoutTensor is already UnsafePointer, so no need to convert to LegacyUnsafePointer at the call sites. The conversion happens inside _store_c_tile where it's actually needed for partial_simd_store.
|
@NathanSWard Hi! I've updated the CPU matmul kernels to use the new The migration keeps CI all green. Good To Merge!! |
…versions - Use free-standing alloc() instead of LegacyUnsafePointer.alloc() - Remove LegacyUnsafePointer conversion in i8mm _store_c_tile - Keep **_ origin syntax for proper type inference - Keep rebind for neon.mojo since _Accumulator uses LegacyUnsafePointer
|
!sync |
|
✅🟣 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. |
|
Landed in 32ae378! Thank you for your contribution 🎉 |
Fixes #5670
This PR migrates the CPU matmul kernel implementations to use the new
UnsafePointerAPI, removing the deprecatedLegacyUnsafePointeralias.Changes
Updated 6 files in
max/kernels/src/linalg/matmul/cpu/:impl.mojo- Core matmul orchestrationapple_accelerate.mojo- Apple Accelerate BLAS wrapperdefault.mojo- Default CPU microkerneli8mm.mojo- ARM i8mm instruction microkernelneon.mojo- ARM NEON microkernelvnni.mojo- x86 VNNI microkernelImplementation Notes
The new
UnsafePointerAPI requires explicit mutability (mut=True/False) and origin tracking. Where needed for compatibility with unmigrated code (like_Accumulator) and FFI interfaces, we userebind[]to convert between pointer types.Testing
bazelw run //bazel/lint:format