mirror of
https://github.com/ml-explore/mlx.git
synced 2025-09-07 17:44:38 +08:00
start cuda circle config (#2256)
* rebase * fix metal kernel linking issue on cuda * start cuda circle config
This commit is contained in:
@@ -17,10 +17,7 @@
|
||||
#include "python/src/indexing.h"
|
||||
#include "python/src/utils.h"
|
||||
|
||||
#include "mlx/device.h"
|
||||
#include "mlx/ops.h"
|
||||
#include "mlx/transforms.h"
|
||||
#include "mlx/utils.h"
|
||||
#include "mlx/mlx.h"
|
||||
|
||||
namespace mx = mlx::core;
|
||||
namespace nb = nanobind;
|
||||
@@ -461,9 +458,12 @@ void init_array(nb::module_& m) {
|
||||
.def(
|
||||
"__dlpack_device__",
|
||||
[](const mx::array& a) {
|
||||
// See
|
||||
// https://github.com/dmlc/dlpack/blob/5c210da409e7f1e51ddf445134a4376fdbd70d7d/include/dlpack/dlpack.h#L74
|
||||
if (mx::metal::is_available()) {
|
||||
// Metal device is available
|
||||
return nb::make_tuple(8, 0);
|
||||
} else if (mx::cu::is_available()) {
|
||||
return nb::make_tuple(13, 0);
|
||||
} else {
|
||||
// CPU device
|
||||
return nb::make_tuple(1, 0);
|
||||
|
@@ -58,4 +58,9 @@ void init_device(nb::module_& m) {
|
||||
&mx::set_default_device,
|
||||
"device"_a,
|
||||
R"pbdoc(Set the default device.)pbdoc");
|
||||
m.def(
|
||||
"is_available",
|
||||
&mx::is_available,
|
||||
"device"_a,
|
||||
R"pbdoc(Check if a back-end is available for the given device.)pbdoc");
|
||||
}
|
||||
|
Reference in New Issue
Block a user