wehub-resource-sync
141 lines
4.6 KiB
C++
141 lines
4.6 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
/*
|
|
Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
|
|
*/
|
|
|
|
#include <cmath>
|
|
#include <iostream>
|
|
|
|
#include "deepspeed_aio_utils.h"
|
|
|
|
using namespace std;
|
|
|
|
const int c_block_size = 128 * 1024;
|
|
const int c_io_queue_depth = 8;
|
|
|
|
io_xfer_ctxt::io_xfer_ctxt(const int fd,
|
|
const int64_t file_offset,
|
|
const int64_t buffer_offset,
|
|
const int64_t num_bytes,
|
|
const void* buffer)
|
|
: _fd(fd),
|
|
_file_base_offset(file_offset),
|
|
_buffer_base_offset(buffer_offset),
|
|
_mem_buffer(buffer),
|
|
_num_bytes(num_bytes)
|
|
{
|
|
}
|
|
|
|
io_prep_context::io_prep_context(const bool read_op,
|
|
const std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
|
|
const size_t block_size,
|
|
const std::vector<struct iocb*>* iocbs)
|
|
: _read_op(read_op), _xfer_ctxt(xfer_ctxt), _block_size(block_size), _iocbs(iocbs)
|
|
{
|
|
}
|
|
|
|
void io_prep_context::prep_iocbs(const int n_iocbs,
|
|
const size_t num_bytes,
|
|
const void* start_buffer,
|
|
const int64_t start_offset)
|
|
{
|
|
assert(static_cast<size_t>(n_iocbs) <= _iocbs->size());
|
|
for (auto i = 0; i < n_iocbs; ++i) {
|
|
const auto shift = i * _block_size;
|
|
const auto xfer_buffer = (char*)start_buffer + _xfer_ctxt->_buffer_base_offset + shift;
|
|
const auto xfer_offset = _xfer_ctxt->_file_base_offset + start_offset + shift;
|
|
auto byte_count = _block_size;
|
|
|
|
if ((shift + _block_size) > num_bytes) { byte_count = num_bytes - shift; }
|
|
|
|
if (_read_op) {
|
|
io_prep_pread(_iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, byte_count, xfer_offset);
|
|
} else {
|
|
io_prep_pwrite(_iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, byte_count, xfer_offset);
|
|
}
|
|
}
|
|
}
|
|
|
|
io_prep_generator::io_prep_generator(const bool read_op,
|
|
const std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
|
|
const size_t block_size)
|
|
: _read_op(read_op),
|
|
_xfer_ctxt(xfer_ctxt),
|
|
_block_size(block_size),
|
|
_remaining_bytes(xfer_ctxt->_num_bytes),
|
|
_next_iocb_index(0)
|
|
{
|
|
_num_io_blocks =
|
|
static_cast<int64_t>(ceil(static_cast<double>(xfer_ctxt->_num_bytes) / block_size));
|
|
_remaining_io_blocks = _num_io_blocks;
|
|
}
|
|
|
|
int io_prep_generator::prep_iocbs(const int n_iocbs, std::vector<struct iocb*>* iocbs)
|
|
{
|
|
if ((_remaining_bytes) == 0 || (_remaining_io_blocks == 0)) {
|
|
assert(static_cast<int64_t>(_remaining_bytes) == _remaining_io_blocks);
|
|
return 0;
|
|
}
|
|
|
|
assert(static_cast<size_t>(n_iocbs) <= iocbs->size());
|
|
|
|
auto actual_n_iocbs = min(static_cast<int64_t>(n_iocbs), _remaining_io_blocks);
|
|
for (auto i = 0; i < actual_n_iocbs; ++i, ++_next_iocb_index) {
|
|
const auto xfer_buffer = (char*)_xfer_ctxt->_mem_buffer + _xfer_ctxt->_buffer_base_offset +
|
|
(_next_iocb_index * _block_size);
|
|
const auto xfer_offset = _xfer_ctxt->_file_base_offset + (_next_iocb_index * _block_size);
|
|
const auto num_bytes = min(static_cast<int64_t>(_block_size), _remaining_bytes);
|
|
if (_read_op) {
|
|
io_prep_pread(iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, num_bytes, xfer_offset);
|
|
} else {
|
|
io_prep_pwrite(iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, num_bytes, xfer_offset);
|
|
}
|
|
_remaining_bytes -= num_bytes;
|
|
}
|
|
_remaining_io_blocks -= actual_n_iocbs;
|
|
|
|
return actual_n_iocbs;
|
|
}
|
|
|
|
int64_t get_file_size(const char* filename, int64_t& size)
|
|
{
|
|
struct stat st;
|
|
if (stat(filename, &st) == -1) { return -1; }
|
|
size = st.st_size;
|
|
return 0;
|
|
}
|
|
|
|
int64_t get_fd_file_size(const int fd, int64_t& size)
|
|
{
|
|
struct stat st;
|
|
if (fstat(fd, &st) == -1) { return -1; }
|
|
size = st.st_size;
|
|
return 0;
|
|
}
|
|
|
|
void* ds_page_aligned_alloc(const int64_t size, const bool lock)
|
|
{
|
|
void* ptr;
|
|
int retval;
|
|
|
|
retval = posix_memalign(&ptr, (size_t)sysconf(_SC_PAGESIZE), size);
|
|
if (retval) { return nullptr; }
|
|
|
|
if (lock == false) { return ptr; }
|
|
|
|
auto mlock_ret = mlock(ptr, size);
|
|
if (mlock_ret != 0) {
|
|
auto mlock_error = errno;
|
|
std::cerr << "mlock failed to allocate " << size << " bytes with error no " << mlock_error
|
|
<< " msg " << strerror(mlock_error) << std::endl;
|
|
free(ptr);
|
|
return nullptr;
|
|
}
|
|
|
|
return ptr;
|
|
}
|