Skip to content

[WIP] adopt new mlx-c api#150

Merged
davidkoski merged 16 commits intoml-explore:mainfrom
davidkoski:mlx-c-new
Dec 4, 2024
Merged

[WIP] adopt new mlx-c api#150
davidkoski merged 16 commits intoml-explore:mainfrom
davidkoski:mlx-c-new

Conversation

@davidkoski
Copy link
Collaborator

@davidkoski davidkoski commented Oct 11, 2024

Adopt new API from ml-explore/mlx-c#38 and move to mlx v0.21.0 (just adopting API, not picking up new ops yet)

This is complete (now that mlx-c is tagged)

@davidkoski davidkoski requested review from andresy and awni October 11, 2024 23:03
@davidkoski davidkoski force-pushed the mlx-c-new branch 2 times, most recently from dcadeaf to b4b2572 Compare November 18, 2024 16:47
let description = mlx_tostring(UnsafeMutableRawPointer(ptr))!
defer { mlx_free(description) }
return String(cString: mlx_string_data(description))
}
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

These no longer take opaque pointers -- the callers can use the mlx_ functions directly and don't need helpers.

// ctx is a +1 object, the array takes ownership
let ctx = mlx_vector_array_get(vector_array, index)!
var ctx = mlx_array_new()
mlx_vector_array_get(&ctx, vector_array, index)
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

99% of the changes in this PR are of this nature: adopt the new API

MLXArray(mlx_zeros(shape.map { Int32($0) }, shape.count, T.dtype.cmlxDtype, stream.ctx))
var result = mlx_array_new()
mlx_zeros(&result, shape.map { Int32($0) }, shape.count, T.dtype.cmlxDtype, stream.ctx)
return MLXArray(result)
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

This is a typical operation:

  • create the result (mlx_array_new())
  • call the function (mlx_zeros)
  • return the result, transfer ownership to it (MLXArray(result))

This comes from https://developer.apple.com/metal/cpp/ specifically:

- https://developer.apple.com/metal/cpp/files/metal-cpp_macOS14.2_iOS17.2.zip
- https://developer.apple.com/metal/cpp/files/metal-cpp_macOS15_iOS18-beta.zip
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Picked up the current version of metal-cpp (per mlx requirements)

@awni
Copy link
Member

awni commented Nov 23, 2024

I'm a bit daunted by the size of the diff 😅 . I can peruse but (other than you inline comments) is there anything specific that would be good to have reviewed?

@davidkoski
Copy link
Collaborator Author

I'm a bit daunted by the size of the diff 😅 . I can peruse but (other than you inline comments) is there anything specific that would be good to have reviewed?

It is mostly mechanical -- I was able to do a good chunk of the work with creative search and replace. The metal kernel code is the biggest change because the structure of the back end changed (and it is complicated anyway).

It is big but mostly uninteresting.

Copy link
Member

@awni awni left a comment

Choose a reason for hiding this comment

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

Tests clear, LGTM. Thanks for getting this one done!!

@davidkoski
Copy link
Collaborator Author

Tests clear, LGTM. Thanks for getting this one done!!

Awesome, thanks for looking! I will merge and tag it in the morning.

@davidkoski davidkoski merged commit bafcb33 into ml-explore:main Dec 4, 2024
@davidkoski davidkoski deleted the mlx-c-new branch December 4, 2024 16:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants