Implementation:Deepspeedai DeepSpeed AIO Utils

Knowledge Sources	DeepSpeed
Domains	Async_IO, NVMe_Offload
Last Updated	2026-02-09 00:00 GMT

Overview

Utility classes and functions for preparing I/O control blocks and managing memory alignment in asynchronous I/O operations.

Description

This module provides essential utility components for DeepSpeed's asynchronous I/O system. It includes the io_xfer_ctxt class for encapsulating transfer context information, two I/O preparation strategies (io_prep_context for batch preparation and io_prep_generator for incremental generation), file size query functions, and page-aligned memory allocation with optional memory locking. These utilities abstract the low-level details of preparing libaio iocb structures and ensure proper memory alignment for O_DIRECT I/O operations.

The io_prep_context prepares a full batch of iocbs at once (used in sequential mode), while io_prep_generator yields iocbs incrementally (used in overlap mode). The page-aligned allocator ensures buffers meet O_DIRECT alignment requirements and optionally locks pages in RAM to prevent swapping.

Usage

These utilities are used throughout DeepSpeed AIO implementations whenever preparing I/O operations, querying file sizes, or allocating aligned buffers for direct I/O. They're invoked internally by higher-level I/O handle classes.

Code Reference

Source Location

Repository: DeepSpeed
File: csrc/aio/common/deepspeed_aio_utils.cpp

Signature

class 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);
};

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

    void prep_iocbs(const int n_iocbs,
                    const size_t num_bytes,
                    const void* start_buffer,
                    const int64_t start_offset);
};

class io_prep_generator {
    io_prep_generator(const bool read_op,
                      const std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
                      const size_t block_size);

    int prep_iocbs(const int n_iocbs, std::vector<struct iocb*>* iocbs);
};

int64_t get_file_size(const char* filename, int64_t& size);
int64_t get_fd_file_size(const int fd, int64_t& size);
void* ds_page_aligned_alloc(const int64_t size, const bool lock);

Import

#include "deepspeed_aio_utils.h"

I/O Contract

Inputs

Name	Type	Required	Description
fd	int	Yes	File descriptor for I/O operations
file_offset	int64_t	Yes	Starting offset in file
buffer_offset	int64_t	Yes	Starting offset in buffer
num_bytes	int64_t	Yes	Total bytes to transfer
buffer	const void*	Yes	Memory buffer pointer
read_op	bool	Yes	True for read, false for write
block_size	size_t	Yes	Size of each I/O block
n_iocbs	int	Yes	Number of iocbs to prepare
filename	const char*	Yes	File path for size query
size	int64_t	Yes	Allocation size in bytes
lock	bool	Yes	Whether to lock pages in memory

Outputs

Name	Type	Description
iocbs	iocb structures	Prepared I/O control blocks ready for submission
actual_n_iocbs	int	Number of iocbs actually prepared (may be less than requested)
size	int64_t	File size in bytes (output parameter)
aligned_ptr	void*	Page-aligned and optionally locked memory pointer, or nullptr on failure
return_code	int64_t	0 on success, -1 on error for file size queries

Usage Examples

// Create transfer context
std::unique_ptr<io_xfer_ctxt> xfer_ctxt(
    new io_xfer_ctxt(fd, 0, 0, num_bytes, buffer));

// Batch preparation for sequential I/O
std::vector<struct iocb*> iocbs;
io_prep_context prep_ctxt(true, xfer_ctxt, 1024*1024, &iocbs);
prep_ctxt.prep_iocbs(32, num_bytes, buffer, 0);

// Incremental generation for overlapped I/O
io_prep_generator io_gen(false, xfer_ctxt, 1024*1024);
while (true) {
    int n = io_gen.prep_iocbs(queue_depth, &iocbs);
    if (n == 0) break;
    // Submit and process iocbs
}

// Get file size
int64_t file_size;
if (get_file_size("/nvme/checkpoint.pt", file_size) == 0) {
    std::cout << "File size: " << file_size << " bytes" << std::endl;
}

// Allocate page-aligned, locked memory
void* buffer = ds_page_aligned_alloc(1024*1024*1024, true);
if (buffer) {
    // Use buffer for O_DIRECT I/O
    munlock(buffer, 1024*1024*1024);
    free(buffer);
}

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment