zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-17 22:51:19 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Remove static initializers in device.cpp, load.cpp, pocketfft.h

Browse Source
This commit is contained in:
Hanyu Deng 2025-04-10 11:43:51 +08:00 committed by Awni Hannun
parent 3836445241
commit aaf261d128
4 changed files with 23 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
mlx/3rdparty/pocketfft.h vendored
View File

@ -537,7 +537,11 @@ inline size_t &num_threads()
static thread_local size_t num_threads_=1; static thread_local size_t num_threads_=1;
return num_threads_; return num_threads_;
} }
static const size_t max_threads = std::max(1u, std::thread::hardware_concurrency()); inline const size_t &max_threads()
{
static thread_local const size_t max_threads_ = std::max(1u, std::thread::hardware_concurrency());
return max_threads_;
}
class latch class latch
{ {
@ -721,7 +725,7 @@ class thread_pool
workers_(nthreads) workers_(nthreads)
{ create_threads(); } { create_threads(); }
thread_pool(): thread_pool(max_threads) {} thread_pool(): thread_pool(max_threads()) {}
~thread_pool() { shutdown(); } ~thread_pool() { shutdown(); }
@ -786,7 +790,7 @@ template <typename Func>
void thread_map(size_t nthreads, Func f) void thread_map(size_t nthreads, Func f)
{ {
if (nthreads == 0) if (nthreads == 0)
nthreads = max_threads; nthreads = max_threads();
if (nthreads == 1) if (nthreads == 1)
{ f(); return; } { f(); return; }

9
mlx/device.cpp
View File

@ -5,11 +5,14 @@
namespace mlx::core { namespace mlx::core {
static Device default_device_{ Device& mutable_default_device() {
static Device default_device{
metal::is_available() ? Device::gpu : Device::cpu}; metal::is_available() ? Device::gpu : Device::cpu};
return default_device;
}
const Device& default_device() { const Device& default_device() {
return default_device_; return mutable_default_device();
} }
void set_default_device(const Device& d) { void set_default_device(const Device& d) {
@ -17,7 +20,7 @@ void set_default_device(const Device& d) {
throw std::invalid_argument( throw std::invalid_argument(
"[set_default_device] Cannot set gpu device without gpu backend."); "[set_default_device] Cannot set gpu device without gpu backend.");
} }
default_device_ = d; mutable_default_device() = d;
} }
bool operator==(const Device& lhs, const Device& rhs) { bool operator==(const Device& lhs, const Device& rhs) {

8
mlx/io/load.cpp
View File

@ -335,7 +335,10 @@ ThreadPool& thread_pool() {
return pool_; return pool_;
} }
ThreadPool ParallelFileReader::thread_pool_{4}; ThreadPool& ParallelFileReader::thread_pool() {
static ThreadPool thread_pool{4};
return thread_pool;
}
void ParallelFileReader::read(char* data, size_t n) { void ParallelFileReader::read(char* data, size_t n) {
while (n != 0) { while (n != 0) {
@ -371,7 +374,8 @@ void ParallelFileReader::read(char* data, size_t n, size_t offset) {
break; break;
} else { } else {
size_t m = batch_size_; size_t m = batch_size_;
futs.emplace_back(thread_pool_.enqueue(readfn, offset, m, data)); futs.emplace_back(
ParallelFileReader::thread_pool().enqueue(readfn, offset, m, data));
data += m; data += m;
n -= m; n -= m;
offset += m; offset += m;

4
mlx/io/load.h
View File

@ -99,9 +99,11 @@ class ParallelFileReader : public Reader {
return "file " + label_; return "file " + label_;
} }
private:
static ThreadPool& thread_pool();
private: private:
static constexpr size_t batch_size_ = 1 << 25; static constexpr size_t batch_size_ = 1 << 25;
static ThreadPool thread_pool_;
int fd_; int fd_;
std::string label_; std::string label_;
}; };