Files
2026-07-13 13:18:33 +08:00

39 lines
885 B
C++

// Copyright (c) Microsoft Corporation.
// SPDX-License-Identifier: Apache-2.0
// DeepSpeed Team
/*
Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
*/
#include <condition_variable>
#include <memory>
#include <queue>
#include "deepspeed_cpu_op.h"
struct thread_sync_t {
std::mutex _mutex;
std::condition_variable _cond_var;
};
struct deepspeed_aio_thread_t {
const int _tid;
deepspeed_aio_config_t& _aio_config;
std::unique_ptr<struct aio_context> _aio_ctxt;
std::queue<std::shared_ptr<struct io_op_desc_t>> _work_queue;
std::queue<std::shared_ptr<struct io_op_desc_t>> _complete_queue;
bool _time_to_exit;
struct thread_sync_t _work_sync;
struct thread_sync_t _complete_sync;
deepspeed_aio_thread_t(const int tid, deepspeed_aio_config_t& aio_config);
~deepspeed_aio_thread_t();
void run();
};