Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added mlx_device_get_index, and mlx_stream_get_index #49

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Added mlx_device_get_index, and mlx_stream_get_index #49

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d82a1e9
Added mlx_device_get_index
stemann Dec 15, 2024
028977e
Also added mlx_stream_get_index
stemann Dec 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions mlx/c/device.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,15 @@ extern "C" int mlx_device_set(mlx_device* dev, const mlx_device src) {
return 0;
}

extern "C" int mlx_device_get_index(mlx_device dev) {
try {
return mlx_device_get_(dev).index;
} catch (std::exception& e) {
mlx_error(e.what());
return 0; // DEBUG: could have a specific value
}
}

extern "C" mlx_device_type mlx_device_get_type(mlx_device dev) {
try {
return mlx_device_type_to_c(mlx_device_get_(dev).type);
Expand All @@ -51,9 +60,11 @@ extern "C" mlx_device_type mlx_device_get_type(mlx_device dev) {
return MLX_CPU; // DEBUG: could have a specific value
}
}

extern "C" bool mlx_device_equal(mlx_device lhs, mlx_device rhs) {
return mlx_device_get_(lhs) == mlx_device_get_(rhs);
}

extern "C" int mlx_get_default_device(mlx_device* dev) {
try {
mlx_device_set_(*dev, mlx::core::default_device());
Expand All @@ -63,6 +74,7 @@ extern "C" int mlx_get_default_device(mlx_device* dev) {
return 1;
}
}

extern "C" int mlx_set_default_device(mlx_device dev) {
try {
mlx::core::set_default_device(mlx_device_get_(dev));
Expand All @@ -72,6 +84,7 @@ extern "C" int mlx_set_default_device(mlx_device dev) {
}
return 0;
}

extern "C" int mlx_device_free(mlx_device dev) {
try {
mlx_device_free_(dev);
Expand Down
8 changes: 8 additions & 0 deletions mlx/c/device.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,14 @@ int mlx_device_set(mlx_device* dev, const mlx_device src);
* Get device description.
*/
int mlx_device_tostring(mlx_string* str, mlx_device dev);
/**
* Check if devices are the same.
*/
bool mlx_device_equal(mlx_device lhs, mlx_device rhs);
/**
* Returns the index of the device.
*/
int mlx_device_get_index(mlx_device dev);
/**
* Returns the type of the device.
*/
Expand Down
8 changes: 8 additions & 0 deletions mlx/c/stream.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,14 @@ extern "C" int mlx_stream_get_device(mlx_device* dev, mlx_stream stream) {
return 1;
}
}
extern "C" int mlx_stream_get_index(mlx_stream stream) {
try {
return mlx_stream_get_(stream).index;
} catch (std::exception& e) {
mlx_error(e.what());
return 0; // DEBUG: could have a specific value
}
}
extern "C" int mlx_synchronize(mlx_stream stream) {
try {
mlx::core::synchronize(mlx_stream_get_(stream));
Expand Down
4 changes: 4 additions & 0 deletions mlx/c/stream.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,10 @@ bool mlx_stream_equal(mlx_stream lhs, mlx_stream rhs);
* Return the device of the stream.
*/
int mlx_stream_get_device(mlx_device* dev, mlx_stream stream);
/**
* Return the index of the stream.
*/
int mlx_stream_get_index(mlx_stream stream);
/**
* Synchronize with the provided stream.
*/
Expand Down