Skip to content

Commit 5bb55ef

Browse files
angeloskathangeloskath
committed
Comments
1 parent 675ee55 commit 5bb55ef

File tree

7 files changed

+49
-51
lines changed

7 files changed

+49
-51
lines changed

‎mlx/backend/cuda/device_info.cpp‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -227,11 +227,6 @@ bool is_available() {
227227
return true;
228228
}
229229

230-
const std::unordered_map<std::string, std::variant<std::string, size_t>>&
231-
device_info(int device_index) {
232-
return device_info_impl(device_index);
233-
}
234-
235230
} // namespace cu
236231

237232
} // namespace mlx::core

‎mlx/backend/cuda/eval.cpp‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
#include "mlx/backend/gpu/eval.h"
44
#include "mlx/backend/cuda/allocator.h"
55
#include "mlx/backend/cuda/device.h"
6-
#include "mlx/backend/gpu/device_info.h"
76
#include "mlx/primitives.h"
87
#include "mlx/scheduler.h"
98

‎mlx/backend/metal/CMakeLists.txt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -139,5 +139,5 @@ endif()
139139

140140
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/kernels)
141141

142-
target_compile_definitions(
143-
mlx PRIVATE METAL_PATH="${MLX_METAL_PATH}/mlx.metallib" MLX_USE_METAL)
142+
target_compile_definitions(mlx
143+
PRIVATE METAL_PATH="${MLX_METAL_PATH}/mlx.metallib")

‎mlx/backend/metal/allocator.cpp‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// Copyright © 2023-2024 Apple Inc.
22
#include "mlx/backend/metal/allocator.h"
3+
#include "mlx/backend/gpu/device_info.h"
34
#include "mlx/backend/metal/metal.h"
45
#include "mlx/backend/metal/resident.h"
56
#include "mlx/memory.h"
@@ -43,16 +44,17 @@ MetalAllocator::MetalAllocator()
4344
}),
4445
residency_set_(device_) {
4546
auto pool = metal::new_scoped_memory_pool();
46-
auto memsize = std::get<size_t>(device_info().at("memory_size"));
47+
const auto& info = gpu::device_info(0);
48+
auto memsize = std::get<size_t>(info.at("memory_size"));
4749
auto max_rec_size =
48-
std::get<size_t>(device_info().at("max_recommended_working_set_size"));
49-
resource_limit_ = std::get<size_t>(device_info().at("resource_limit"));
50+
std::get<size_t>(info.at("max_recommended_working_set_size"));
51+
resource_limit_ = std::get<size_t>(info.at("resource_limit"));
5052
block_limit_ = std::min(1.5 * max_rec_size, 0.95 * memsize);
5153
gc_limit_ = std::min(static_cast<size_t>(0.95 * max_rec_size), block_limit_);
5254
max_pool_size_ = block_limit_;
5355
device(mlx::core::Device::gpu)
5456
.set_residency_set(residency_set_.mtl_residency_set());
55-
bool is_vm = std::get<std::string>(device_info().at("device_name")) ==
57+
bool is_vm = std::get<std::string>(info.at("device_name")) ==
5658
"Apple Paravirtual device";
5759
if (is_vm) {
5860
return;
@@ -249,8 +251,8 @@ size_t get_memory_limit() {
249251
return metal::allocator().get_memory_limit();
250252
}
251253
size_t set_wired_limit(size_t limit) {
252-
if (limit > std::get<size_t>(metal::device_info().at(
253-
"max_recommended_working_set_size"))) {
254+
if (limit > std::get<size_t>(
255+
gpu::device_info(0).at("max_recommended_working_set_size"))) {
254256
throw std::invalid_argument(
255257
"[metal::set_wired_limit] Setting a wired limit larger than "
256258
"the maximum working set size is not allowed.");

‎mlx/backend/metal/device_info.cpp‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
// Copyright © 2026 Apple Inc.
22

3+
#include <sys/sysctl.h>
4+
35
#include "mlx/backend/gpu/device_info.h"
6+
#include "mlx/backend/metal/device.h"
47
#include "mlx/backend/metal/metal.h"
58

69
namespace mlx::core::gpu {
@@ -14,8 +17,42 @@ int device_count() {
1417
}
1518

1619
const std::unordered_map<std::string, std::variant<std::string, size_t>>&
17-
device_info(int /* device_index */) {
18-
return metal::device_info();
20+
device_info(int device_index) {
21+
auto init_device_info = []()
22+
-> std::unordered_map<std::string, std::variant<std::string, size_t>> {
23+
auto pool = metal::new_scoped_memory_pool();
24+
auto raw_device = metal::device(mlx::core::Device::gpu).mtl_device();
25+
auto name = std::string(raw_device->name()->utf8String());
26+
auto arch = std::string(raw_device->architecture()->name()->utf8String());
27+
28+
size_t memsize = 0;
29+
size_t length = sizeof(memsize);
30+
sysctlbyname("hw.memsize", &memsize, &length, NULL, 0);
31+
32+
size_t rsrc_limit = 0;
33+
sysctlbyname("iogpu.rsrc_limit", &rsrc_limit, &length, NULL, 0);
34+
if (rsrc_limit == 0) {
35+
rsrc_limit = 499000;
36+
}
37+
38+
return {
39+
{"device_name", name},
40+
{"architecture", arch},
41+
{"max_buffer_length", raw_device->maxBufferLength()},
42+
{"max_recommended_working_set_size",
43+
raw_device->recommendedMaxWorkingSetSize()},
44+
{"memory_size", memsize},
45+
{"resource_limit", rsrc_limit}};
46+
};
47+
static auto device_info_ = init_device_info();
48+
static std::unordered_map<std::string, std::variant<std::string, size_t>>
49+
empty;
50+
51+
if (device_index == 0) {
52+
return device_info_;
53+
} else {
54+
return empty;
55+
}
1956
}
2057

2158
} // namespace mlx::core::gpu

‎mlx/backend/metal/eval.cpp‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
// Copyright © 2023-2024 Apple Inc.
22
#include <memory>
33

4-
#include "mlx/backend/gpu/device_info.h"
54
#include "mlx/backend/gpu/eval.h"
65
#include "mlx/backend/metal/device.h"
76
#include "mlx/backend/metal/utils.h"

‎mlx/backend/metal/metal.cpp‎

Lines changed: 0 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
// Copyright © 2023-2024 Apple Inc.
22
#include <memory>
33

4-
#include <sys/sysctl.h>
5-
64
#include "mlx/backend/metal/device.h"
75
#include "mlx/backend/metal/metal.h"
86
#include "mlx/backend/metal/utils.h"
@@ -49,36 +47,4 @@ void stop_capture() {
4947
manager->stopCapture();
5048
}
5149

52-
const std::unordered_map<std::string, std::variant<std::string, size_t>>&
53-
device_info() {
54-
auto init_device_info = []()
55-
-> std::unordered_map<std::string, std::variant<std::string, size_t>> {
56-
auto pool = new_scoped_memory_pool();
57-
auto raw_device = device(default_device()).mtl_device();
58-
auto name = std::string(raw_device->name()->utf8String());
59-
auto arch = std::string(raw_device->architecture()->name()->utf8String());
60-
61-
size_t memsize = 0;
62-
size_t length = sizeof(memsize);
63-
sysctlbyname("hw.memsize", &memsize, &length, NULL, 0);
64-
65-
size_t rsrc_limit = 0;
66-
sysctlbyname("iogpu.rsrc_limit", &rsrc_limit, &length, NULL, 0);
67-
if (rsrc_limit == 0) {
68-
rsrc_limit = 499000;
69-
}
70-
71-
return {
72-
{"device_name", name},
73-
{"architecture", arch},
74-
{"max_buffer_length", raw_device->maxBufferLength()},
75-
{"max_recommended_working_set_size",
76-
raw_device->recommendedMaxWorkingSetSize()},
77-
{"memory_size", memsize},
78-
{"resource_limit", rsrc_limit}};
79-
};
80-
static auto device_info_ = init_device_info();
81-
return device_info_;
82-
}
83-
8450
} // namespace mlx::core::metal

0 commit comments

Comments
 (0)