start cuda circle config (#2256)

* rebase

* fix metal kernel linking issue on cuda

* start cuda circle config
This commit is contained in:
Awni Hannun
2025-06-10 21:19:47 -07:00
committed by GitHub
parent 8590c0941e
commit c35f4d089a
14 changed files with 101 additions and 26 deletions

View File

@@ -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);

View File

@@ -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");
}