diff --git a/src/ailego/io/io_backend_def.h b/src/ailego/io/io_backend_def.h index d795c3046..6826c8bca 100644 --- a/src/ailego/io/io_backend_def.h +++ b/src/ailego/io/io_backend_def.h @@ -12,16 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Abstract I/O backend selector. +// Abstract I/O backend selector — internal header. // -// Wraps the low-level loaders (LibAioLoader for libaio) and provides a uniform -// way to initialize, query, and report the active I/O backend. The actual I/O -// operations are still performed by the underlying loaders; this class is -// responsible only for backend initialization and reporting. +// Wraps the low-level backends (io_uring via raw syscalls, LibAioLoader for +// libaio) and provides a uniform way to initialize, query, and report the +// active I/O backend. The actual I/O operations are still performed by the +// underlying backends; this class is responsible only for backend +// initialization and reporting. // // When no async backend is available, the caller should fall back to // synchronous pread(). // +// This header pulls in libaio_loader and the io_uring kernel ABI; the +// dependency-free enum and IOBackendTypeName() live in the public header +// zvec/ailego/io/io_backend.h, which this header includes. +// // Usage: // auto& backend = ailego::IOBackend::Instance(); // if (!backend.is_pread()) { ... } @@ -32,24 +37,25 @@ #include #include +#if defined(__linux) || defined(__linux__) +#include // ::syscall(), ::close() — POSIX only +#include // std::memset +#include // io_uring_params, __NR_io_uring_setup +#endif + namespace zvec { namespace ailego { -// Returns a human-readable name for the given backend type. -inline const char *IOBackendTypeName(IOBackendType type) { - switch (type) { - case IOBackendType::kLibAio: - return "libaio"; - case IOBackendType::kPread: - return "pread"; - } - return "unknown"; -} +// IOBackendTypeName() is defined in the public header +// zvec/ailego/io/io_backend.h. // Returns a human-readable description for the given backend type. // When the backend is kPread, includes installation guidance for libaio. inline const char *IOBackendDescription(IOBackendType type) { switch (type) { + case IOBackendType::kIoUring: + return "io_uring async I/O backend (raw kernel syscalls, zero " + "dependency)."; case IOBackendType::kLibAio: return "libaio async I/O backend loaded at runtime via dlopen()."; case IOBackendType::kPread: @@ -63,8 +69,8 @@ inline const char *IOBackendDescription(IOBackendType type) { // Singleton that loads and queries an I/O backend on demand. // -// available() (no arg) tries the best backend with priority (libaio > pread) -// and returns the loaded backend type. +// available() (no arg) tries the best backend with priority +// (io_uring > libaio > pread) and returns the loaded backend type. // available(IOBackendType) tries a specific backend. // Use type() / name() to query the loaded backend without triggering a load. class IOBackend { @@ -74,14 +80,18 @@ class IOBackend { return instance; } - // Try to load the best available backend (libaio > pread). + // Try to load the best available backend (io_uring > libaio > pread). // Returns the loaded backend type. // Idempotent — if already loaded, returns immediately. IOBackendType available() { if (type_ != IOBackendType::kPread) { return type_; } - return available(IOBackendType::kLibAio); + IOBackendType t = available(IOBackendType::kIoUring); + if (t == IOBackendType::kPread) { + t = available(IOBackendType::kLibAio); + } + return t; } // Try to load the requested backend. Returns the loaded backend type @@ -92,6 +102,18 @@ class IOBackend { return type_; } #if defined(__linux) || defined(__linux__) + if (requested == IOBackendType::kIoUring) { + // Probe io_uring availability with a minimal ring setup using only + // raw syscalls — no dependency on liburing. + struct io_uring_params params; + std::memset(¶ms, 0, sizeof(params)); + int fd = static_cast(::syscall(__NR_io_uring_setup, 1, ¶ms)); + if (fd >= 0) { + ::close(fd); + type_ = IOBackendType::kIoUring; + return type_; + } + } if (requested == IOBackendType::kLibAio) { if (LibAioLoader::Instance().load() && LibAioLoader::Instance().is_available()) { @@ -112,6 +134,10 @@ class IOBackend { return available() == IOBackendType::kLibAio; } + bool is_io_uring() { + return available() == IOBackendType::kIoUring; + } + // Returns the loaded backend type. IOBackendType type() const { return type_; diff --git a/src/ailego/io/iouring_def.h b/src/ailego/io/iouring_def.h new file mode 100644 index 000000000..ce87b246e --- /dev/null +++ b/src/ailego/io/iouring_def.h @@ -0,0 +1,199 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Private header defining the io_uring kernel ABI structures and constants. +// +// This header is the io_uring analogue of libaio_def.h: it declares *only* +// the types, constants, and inline helpers that zvec needs from the io_uring +// kernel interface. By defining these structures ourselves we avoid any +// build-time dependency on or liburing-dev, mirroring the +// project's zero-dependency philosophy established by the libaio dlopen +// approach. +// +// The struct layouts (io_uring_sqe, io_uring_cqe, io_uring_params, +// io_sqring_offsets, io_cqring_offsets) are part of the Linux kernel ABI +// and are copied verbatim from . + +#pragma once + +#include + +#if defined(__linux) || defined(__linux__) + +// --------------------------------------------------------------------------- +// Syscall numbers +// --------------------------------------------------------------------------- +// io_uring was introduced in Linux 5.1 (2019). The three syscalls share the +// same numbers across all supported architectures. We prefer the values +// from when available and fall back to hardcoded numbers. +#include + +#ifndef __NR_io_uring_setup +#define __NR_io_uring_setup 425 +#endif + +#ifndef __NR_io_uring_enter +#define __NR_io_uring_enter 426 +#endif + +#ifndef __NR_io_uring_register +#define __NR_io_uring_register 427 +#endif + +// --------------------------------------------------------------------------- +// Constants (from ) +// --------------------------------------------------------------------------- + +// mmap offsets for the three shared regions. +#define IORING_OFF_SQ_RING 0ULL +#define IORING_OFF_CQ_RING 0x8000000ULL +#define IORING_OFF_SQES 0x10000000ULL + +// io_uring_enter flags. +#define IORING_ENTER_GETEVENTS (1U << 0) + +// io_uring_setup flags (none used by default). +// IORING_SETUP_IOPOLL (1U << 0) +// IORING_SETUP_SQPOLL (1U << 1) +// IORING_SETUP_SQ_AFF (1U << 2) + +// SQE opcode values. +#define IORING_OP_NOP 0 +#define IORING_OP_READV 1 +#define IORING_OP_WRITEV 2 +#define IORING_OP_FSYNC 3 +#define IORING_OP_READ_FIXED 4 +#define IORING_OP_WRITE_FIXED 5 +#define IORING_OP_POLL_ADD 6 +#define IORING_OP_POLL_REMOVE 7 +#define IORING_OP_SYNC_FILE_RANGE 8 +#define IORING_OP_SENDMSG 9 +#define IORING_OP_RECVMSG 10 +#define IORING_OP_TIMEOUT 11 +#define IORING_OP_TIMEOUT_REMOVE 12 +#define IORING_OP_ACCEPT 13 +#define IORING_OP_ASYNC_CANCEL 14 +#define IORING_OP_LINK_TIMEOUT 15 +#define IORING_OP_CONNECT 16 +#define IORING_OP_FALLOCATE 17 +#define IORING_OP_OPENAT 18 +#define IORING_OP_CLOSE 19 +#define IORING_OP_FILES_UPDATE 20 +#define IORING_OP_STATX 21 +#define IORING_OP_READ 22 +#define IORING_OP_WRITE 23 + +// --------------------------------------------------------------------------- +// Struct definitions (copied verbatim from ) +// --------------------------------------------------------------------------- + +// Submission queue entry — 64 bytes. +struct io_uring_sqe { + uint8_t opcode; // type of operation for this sqe + uint8_t flags; // IOSQE_ flags + uint16_t ioprio; // ioprio for the request + int32_t fd; // file descriptor to do IO on + union { + uint64_t off; // offset into file + uint64_t addr2; + }; + union { + uint64_t addr; // buffer or iovecs + uint64_t splice_off_in; + }; + uint32_t len; // buffer size or number of iovecs + union { + uint32_t rw_flags; // read/write flags (union of all flag types) + }; + uint64_t user_data; // data to be passed back at completion time + union { + struct { + uint16_t buf_index; // index into fixed buffers, if used + uint16_t personality; + } buf; + uint64_t __pad2[3]; + }; +}; + +// Completion queue entry — 16 bytes. +struct io_uring_cqe { + uint64_t user_data; // sqe->user_data + int32_t res; // result code for this event + uint32_t flags; +}; + +// SQ ring offsets — returned by io_uring_setup in io_uring_params. +struct io_sqring_offsets { + uint32_t head; + uint32_t tail; + uint32_t ring_mask; + uint32_t ring_entries; + uint32_t flags; + uint32_t dropped; + uint32_t array; + uint32_t resv1; + uint64_t resv2; +}; + +// CQ ring offsets — returned by io_uring_setup in io_uring_params. +struct io_cqring_offsets { + uint32_t head; + uint32_t tail; + uint32_t ring_mask; + uint32_t ring_entries; + uint32_t overflow; + uint32_t cqes; + uint32_t flags; + uint32_t resv1; + uint64_t resv2; +}; + +// Parameters passed to io_uring_setup(). +struct io_uring_params { + uint32_t sq_entries; + uint32_t cq_entries; + uint32_t flags; + uint32_t sq_thread_cpu; + uint32_t sq_thread_idle; + uint32_t features; + uint32_t wq_fd; + uint32_t resv[3]; + struct io_sqring_offsets sq_off; + struct io_cqring_offsets cq_off; +}; + +// --------------------------------------------------------------------------- +// Inline helper — prepare an SQE for a read operation. +// --------------------------------------------------------------------------- +static inline void io_uring_prep_read(struct io_uring_sqe *sqe, int fd, + void *buf, uint32_t nbytes, + uint64_t offset) { + sqe->opcode = IORING_OP_READ; + sqe->flags = 0; + sqe->ioprio = 0; + sqe->fd = fd; + sqe->off = offset; + sqe->addr = reinterpret_cast(buf); + sqe->len = nbytes; + sqe->rw_flags = 0; + sqe->user_data = 0; + sqe->buf.buf_index = 0; + sqe->buf.personality = 0; +} + +// --------------------------------------------------------------------------- +// End: struct and constant definitions from +// --------------------------------------------------------------------------- + +#endif // __linux__ diff --git a/src/ailego/io/iouring_loader.h b/src/ailego/io/iouring_loader.h new file mode 100644 index 000000000..be997d807 --- /dev/null +++ b/src/ailego/io/iouring_loader.h @@ -0,0 +1,259 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Raw-syscall wrapper for Linux io_uring. +// +// This class implements the io_uring submission/completion queue lifecycle +// using *only* kernel syscalls (io_uring_setup, io_uring_enter) and mmap. +// There is zero dependency on liburing or liburing-dev: no build-time header, +// no runtime .so, and no dlopen. This mirrors the project's existing +// zero-dependency philosophy established by the libaio dlopen approach, but +// goes one step further — io_uring is a pure kernel ABI accessed via syscall. +// +// Runtime detection is automatic: if the kernel does not support io_uring +// (pre-5.1 or io_uring disabled), io_uring_setup() returns -ENOSYS and +// setup() returns false, allowing callers to fall back to libaio or pread. +// +// The ring is **not** thread-safe. Each thread that performs I/O must have +// its own IoUringRing instance (managed through IoBackend / IOContext). + +#pragma once + +#if defined(__linux) || defined(__linux__) + +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +// Forward declaration — AlignedRead is defined in diskann_file_reader.h +// after this header is included. The forward declaration is sufficient +// because IoUringRing::execute takes it by reference. +struct AlignedRead; + +// Maximum number of SQEs we submit in a single io_uring_enter() call. +static constexpr uint32_t kIoUringMaxBatch = 128; + +class IoUringRing { + public: + IoUringRing() = default; + ~IoUringRing() { + teardown(); + } + + IoUringRing(const IoUringRing &) = delete; + IoUringRing &operator=(const IoUringRing &) = delete; + + // Create an io_uring with the given number of entries. + // Returns true on success, false if the kernel does not support io_uring + // or setup failed for any reason. + bool setup(uint32_t entries) { + struct io_uring_params params; + std::memset(¶ms, 0, sizeof(params)); + + // io_uring_setup is a raw syscall — returns fd (>=0) or -1 with errno. + ring_fd_ = static_cast( + syscall(__NR_io_uring_setup, static_cast(entries), ¶ms)); + if (ring_fd_ < 0) { + // ENOSYS = kernel doesn't support io_uring. + // EPERM = io_uring disabled via sysctl. + // EINVAL = invalid parameters. + if (errno != ENOSYS) { + LOG_WARN("io_uring_setup failed; errno=%d, %s", errno, + ::strerror(errno)); + } + return false; + } + + sq_entries_ = params.sq_entries; + cq_entries_ = params.cq_entries; + + // --- mmap the three shared regions --- + + // 1. SQ ring (includes head, tail, mask, entries, flags, dropped, array). + size_t sq_ring_sz = static_cast(params.sq_off.array) + + sq_entries_ * sizeof(uint32_t); + sq_ring_ptr_ = ::mmap(nullptr, sq_ring_sz, PROT_READ | PROT_WRITE, + MAP_SHARED, ring_fd_, IORING_OFF_SQ_RING); + if (sq_ring_ptr_ == MAP_FAILED) { + LOG_ERROR("mmap SQ ring failed: %s", ::strerror(errno)); + sq_ring_ptr_ = nullptr; + teardown(); + return false; + } + + // 2. SQE array. + size_t sqes_sz = sq_entries_ * sizeof(struct io_uring_sqe); + sqes_ptr_ = reinterpret_cast( + ::mmap(nullptr, sqes_sz, PROT_READ | PROT_WRITE, MAP_SHARED, ring_fd_, + IORING_OFF_SQES)); + if (sqes_ptr_ == MAP_FAILED) { + LOG_ERROR("mmap SQEs failed: %s", ::strerror(errno)); + sqes_ptr_ = nullptr; + teardown(); + return false; + } + + // 3. CQ ring (includes head, tail, mask, entries, overflow, cqes[]). + size_t cq_ring_sz = static_cast(params.cq_off.cqes) + + cq_entries_ * sizeof(struct io_uring_cqe); + cq_ring_ptr_ = ::mmap(nullptr, cq_ring_sz, PROT_READ | PROT_WRITE, + MAP_SHARED, ring_fd_, IORING_OFF_CQ_RING); + if (cq_ring_ptr_ == MAP_FAILED) { + LOG_ERROR("mmap CQ ring failed: %s", ::strerror(errno)); + cq_ring_ptr_ = nullptr; + teardown(); + return false; + } + + // --- Set up typed pointers into the mmap'd regions --- + + // SQ ring fields. + sq_head_ = reinterpret_cast(static_cast(sq_ring_ptr_) + + params.sq_off.head); + sq_tail_ = reinterpret_cast(static_cast(sq_ring_ptr_) + + params.sq_off.tail); + sq_ring_mask_ = reinterpret_cast( + static_cast(sq_ring_ptr_) + params.sq_off.ring_mask); + sq_ring_entries_ = reinterpret_cast( + static_cast(sq_ring_ptr_) + params.sq_off.ring_entries); + sq_flags_ = reinterpret_cast(static_cast(sq_ring_ptr_) + + params.sq_off.flags); + sq_dropped_ = reinterpret_cast( + static_cast(sq_ring_ptr_) + params.sq_off.dropped); + sq_array_ = reinterpret_cast(static_cast(sq_ring_ptr_) + + params.sq_off.array); + + // CQ ring fields. + cq_head_ = reinterpret_cast(static_cast(cq_ring_ptr_) + + params.cq_off.head); + cq_tail_ = reinterpret_cast(static_cast(cq_ring_ptr_) + + params.cq_off.tail); + cq_ring_mask_ = reinterpret_cast( + static_cast(cq_ring_ptr_) + params.cq_off.ring_mask); + cq_ring_entries_ = reinterpret_cast( + static_cast(cq_ring_ptr_) + params.cq_off.ring_entries); + cq_overflow_ = reinterpret_cast( + static_cast(cq_ring_ptr_) + params.cq_off.overflow); + cqes_ = reinterpret_cast( + static_cast(cq_ring_ptr_) + params.cq_off.cqes); + + // SQE array. + sqes_ = sqes_ptr_; + + // Initialize the SQ array to identity mapping so that logical index == + // physical SQE index. This is the simplest and most common configuration. + for (uint32_t i = 0; i < sq_entries_; i++) { + sq_array_[i] = i; + } + + return true; + } + + // Tear down the ring: munmap all regions and close the ring fd. + void teardown() { + if (sq_ring_ptr_ && sq_ring_ptr_ != MAP_FAILED) { + // We don't track the exact mmap size; munmap with a large enough size + // is safe because the kernel only unmaps what was actually mapped. + // However, to be correct we use the page-aligned size. + size_t sz = static_cast(sq_entries_) * sizeof(uint32_t) + 4096; + ::munmap(sq_ring_ptr_, sz); + } + if (sqes_ptr_ && sqes_ptr_ != MAP_FAILED) { + size_t sz = + static_cast(sq_entries_) * sizeof(struct io_uring_sqe); + ::munmap(sqes_ptr_, sz); + } + if (cq_ring_ptr_ && cq_ring_ptr_ != MAP_FAILED) { + size_t sz = + static_cast(cq_entries_) * sizeof(struct io_uring_cqe) + 4096; + ::munmap(cq_ring_ptr_, sz); + } + + sq_ring_ptr_ = nullptr; + sqes_ptr_ = nullptr; + cq_ring_ptr_ = nullptr; + sqes_ = nullptr; + cqes_ = nullptr; + sq_head_ = sq_tail_ = sq_ring_mask_ = sq_ring_entries_ = nullptr; + sq_flags_ = sq_dropped_ = sq_array_ = nullptr; + cq_head_ = cq_tail_ = cq_ring_mask_ = cq_ring_entries_ = nullptr; + cq_overflow_ = nullptr; + + if (ring_fd_ >= 0) { + ::close(ring_fd_); + ring_fd_ = -1; + } + + sq_entries_ = 0; + cq_entries_ = 0; + } + + bool is_valid() const { + return ring_fd_ >= 0; + } + + // Execute a batch of aligned read requests via io_uring. + // + // Implemented in diskann_file_reader.cc to avoid a circular dependency: + // AlignedRead is defined in diskann_file_reader.h *after* this header is + // included, so the method body cannot be inline here. + // + // On success returns 0. On failure returns -1; the caller should fall + // back to pread. + int execute(int fd, std::vector &read_reqs); + + private: + int ring_fd_{-1}; + + // mmap'd region base pointers (needed for munmap). + void *sq_ring_ptr_{nullptr}; + struct io_uring_sqe *sqes_ptr_{nullptr}; + void *cq_ring_ptr_{nullptr}; + + // SQ ring field pointers (into sq_ring_ptr_). + unsigned *sq_head_{nullptr}; + unsigned *sq_tail_{nullptr}; + unsigned *sq_ring_mask_{nullptr}; + unsigned *sq_ring_entries_{nullptr}; + unsigned *sq_flags_{nullptr}; + unsigned *sq_dropped_{nullptr}; + unsigned *sq_array_{nullptr}; + + // CQ ring field pointers (into cq_ring_ptr_). + unsigned *cq_head_{nullptr}; + unsigned *cq_tail_{nullptr}; + unsigned *cq_ring_mask_{nullptr}; + unsigned *cq_ring_entries_{nullptr}; + unsigned *cq_overflow_{nullptr}; + struct io_uring_cqe *cqes_{nullptr}; + + // SQE array pointer. + struct io_uring_sqe *sqes_{nullptr}; + + // Ring capacities. + unsigned sq_entries_{0}; + unsigned cq_entries_{0}; +}; + +} // namespace core +} // namespace zvec + +#endif // __linux__ diff --git a/src/binding/python/typing/python_type.cc b/src/binding/python/typing/python_type.cc index d20cac755..45d3d8216 100644 --- a/src/binding/python/typing/python_type.cc +++ b/src/binding/python/typing/python_type.cc @@ -146,6 +146,7 @@ Enumeration of supported I/O backend types for DiskAnn async disk reads. - PREAD: Synchronous pread() \u2014 no async I/O. - LIBAIO: libaio loaded at runtime via dlopen(). +- IO_URING: io_uring via raw kernel syscalls (zero dependency). Examples: >>> from zvec.typing import IOBackendType @@ -153,7 +154,8 @@ Enumeration of supported I/O backend types for DiskAnn async disk reads. IOBackendType.LIBAIO )pbdoc") .value("PREAD", ailego::IOBackendType::kPread) - .value("LIBAIO", ailego::IOBackendType::kLibAio); + .value("LIBAIO", ailego::IOBackendType::kLibAio) + .value("IO_URING", ailego::IOBackendType::kIoUring); } void ZVecPyTyping::bind_status(py::module_ &m) { diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 85be8c334..e26533c4f 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -57,11 +57,37 @@ int setup_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) std::call_once(g_io_backend_log_once, log_diskann_io_backend); if (ailego::IOBackend::Instance().is_pread()) { + // No async backend available — leave ctx null so callers fall back to + // synchronous pread(). return 0; } - int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); - return ret; + ctx = new IoBackend(); + ailego::IOBackendType selected = ailego::IOBackend::Instance().available(); + + // Priority 1: io_uring (raw kernel syscalls — zero dependency). + if (selected == ailego::IOBackendType::kIoUring && + ctx->ring.setup(MAX_EVENTS)) { + ctx->backend = IoBackend::IO_URING; + return 0; + } + + // Priority 2: libaio (dlopen — soft dependency). + if (selected != ailego::IOBackendType::kPread && + LibAioLoader::Instance().load() && + LibAioLoader::Instance().is_available()) { + int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx->aio_ctx); + if (ret == 0) { + ctx->backend = IoBackend::LIBAIO; + return 0; + } + LOG_WARN("io_setup failed; returned: %d, %s. falling back to pread", ret, + ::strerror(-ret)); + } + + // Priority 3: synchronous pread (always available). + ctx->backend = IoBackend::NONE; + return 0; #else return 0; #endif @@ -69,12 +95,21 @@ int setup_io_ctx(IOContext &ctx) { int destroy_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) - if (ailego::IOBackend::Instance().is_pread()) { + if (ctx == nullptr) { return 0; } - int ret = LibAioLoader::Instance().io_destroy(ctx); - return ret; + if (ctx->backend == IoBackend::IO_URING) { + ctx->ring.teardown(); + } else if (ctx->backend == IoBackend::LIBAIO && + LibAioLoader::Instance().is_available()) { + LibAioLoader::Instance().io_destroy(ctx->aio_ctx); + } + // IoUringRing destructor also calls teardown() — idempotent and safe. + + delete ctx; + ctx = nullptr; + return 0; #else return 0; #endif @@ -98,8 +133,9 @@ static int execute_io_pread(int fd, std::vector &read_reqs) { return 0; } -int execute_io(IOContext ctx, int fd, std::vector &read_reqs, - uint64_t n_retries = 0) { +static int execute_io_libaio(io_context_t ctx, int fd, + std::vector &read_reqs, + uint64_t n_retries = 0) { #if (defined(__linux) || defined(__linux__)) if (ailego::IOBackend::Instance().is_pread()) { return execute_io_pread(fd, read_reqs); @@ -183,6 +219,135 @@ int execute_io(IOContext ctx, int fd, std::vector &read_reqs, #endif } +int execute_io(IOContext ctx, int fd, std::vector &read_reqs, + uint64_t n_retries = 0) { +#if (defined(__linux) || defined(__linux__)) + // Guard against null or sentinel contexts. + if (ctx == nullptr || ctx == (IOContext)-1) { + return execute_io_pread(fd, read_reqs); + } + + // Dispatch based on the active backend. + if (ctx->backend == IoBackend::IO_URING) { + int ret = ctx->ring.execute(fd, read_reqs); + if (ret == 0) { + return 0; + } + // io_uring failed — fall back to pread. + LOG_WARN("io_uring execute failed; falling back to pread"); + return execute_io_pread(fd, read_reqs); + } + + if (ctx->backend == IoBackend::LIBAIO) { + return execute_io_libaio(ctx->aio_ctx, fd, read_reqs, n_retries); + } + + // NONE backend — synchronous pread. + return execute_io_pread(fd, read_reqs); +#else + return execute_io_pread(fd, read_reqs); +#endif +} + +// --------------------------------------------------------------------------- +// IoUringRing::execute — defined here (not in iouring_loader.h) because it +// accesses AlignedRead members, and AlignedRead is defined in +// diskann_file_reader.h after iouring_loader.h is included. +// --------------------------------------------------------------------------- +#if (defined(__linux) || defined(__linux__)) +int IoUringRing::execute(int fd, std::vector &read_reqs) { + if (!is_valid()) { + return -1; + } + if (read_reqs.empty()) { + return 0; + } + + // Process in batches limited by the SQ ring size. + uint32_t batch_size = + std::min(sq_entries_, static_cast(kIoUringMaxBatch)); + uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), batch_size); + + for (uint64_t iter = 0; iter < iters; iter++) { + uint64_t n_ops = + std::min(static_cast(read_reqs.size()) - iter * batch_size, + static_cast(batch_size)); + + // --- Phase 1: Fill SQEs --- + + unsigned tail = __atomic_load_n(sq_tail_, __ATOMIC_ACQUIRE); + unsigned mask = *sq_ring_mask_; + + for (uint64_t j = 0; j < n_ops; j++) { + unsigned idx = (tail + static_cast(j)) & mask; + unsigned sqe_idx = sq_array_[idx]; + struct io_uring_sqe *sqe = &sqes_[sqe_idx]; + + uint64_t req_idx = j + iter * batch_size; + io_uring_prep_read(sqe, fd, read_reqs[req_idx].buf, + static_cast(read_reqs[req_idx].len), + read_reqs[req_idx].offset); + // Store the request index so we can verify the completion. + sqe->user_data = req_idx; + } + + // Memory barrier: ensure SQE contents are visible before tail update. + __sync_synchronize(); + __atomic_store_n(sq_tail_, tail + static_cast(n_ops), + __ATOMIC_RELEASE); + + // --- Phase 2: Submit and wait for completions --- + + int ret = static_cast( + syscall(__NR_io_uring_enter, ring_fd_, static_cast(n_ops), + static_cast(n_ops), IORING_ENTER_GETEVENTS, + static_cast(nullptr), static_cast(0))); + if (ret < 0) { + LOG_WARN( + "io_uring_enter failed; errno=%d, %s, n_ops=%lu. " + "falling back to pread", + errno, ::strerror(errno), (unsigned long)n_ops); + return -1; + } + + // --- Phase 3: Process CQEs --- + + unsigned head = __atomic_load_n(cq_head_, __ATOMIC_ACQUIRE); + unsigned cq_mask = *cq_ring_mask_; + bool all_ok = true; + uint64_t processed = 0; + + for (unsigned i = head; processed < n_ops; i = (i + 1), processed++) { + struct io_uring_cqe *cqe = &cqes_[i & cq_mask]; + uint64_t req_idx = cqe->user_data; + + if (cqe->res < 0) { + LOG_WARN("io_uring read failed: req=%lu, res=%d, offset=%lu", + (unsigned long)req_idx, cqe->res, + (unsigned long)read_reqs[req_idx].offset); + all_ok = false; + } else if (static_cast(cqe->res) != read_reqs[req_idx].len) { + LOG_WARN("io_uring short read: req=%lu, got=%d, expected=%lu", + (unsigned long)req_idx, cqe->res, + (unsigned long)read_reqs[req_idx].len); + all_ok = false; + } + } + + // Advance the CQ head to consume the completions. + __sync_synchronize(); + __atomic_store_n(cq_head_, head + static_cast(n_ops), + __ATOMIC_RELEASE); + + if (!all_ok) { + return -1; + } + } + + return 0; +} +#endif // __linux__ + LinuxAlignedFileReader::LinuxAlignedFileReader(int file_desc) { this->file_desc = file_desc; } @@ -203,7 +368,7 @@ IOContext &LinuxAlignedFileReader::get_ctx() { std::unique_lock lk(ctx_mut); auto it = ctx_map.find(std::this_thread::get_id()); if (it == ctx_map.end()) { - LOG_ERROR("bad thread access; returning -1 as io_context_t"); + LOG_ERROR("bad thread access; returning invalid IOContext"); return this->bad_ctx; } else { return it->second; @@ -216,32 +381,20 @@ void LinuxAlignedFileReader::register_thread() { std::unique_lock lk(ctx_mut); if (ctx_map.find(thread_id) != ctx_map.end()) { LOG_ERROR("multiple calls to register_thread from the same thread"); - return; } IOContext ctx = nullptr; - - std::call_once(g_io_backend_log_once, log_diskann_io_backend); - if (ailego::IOBackend::Instance().is_pread()) { + int ret = setup_io_ctx(ctx); + if (ret != 0) { + LOG_ERROR("setup_io_ctx failed; returned: %d", ret); lk.unlock(); return; } - int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); - if (ret != 0) { - if (ret == -EAGAIN) { - LOG_ERROR( - "io_setup failed with EAGAIN: Consider increasing " - "/proc/sys/fs/aio-max-nr"); - } else { - LOG_ERROR("io_setup failed; returned: %d, %s", ret, ::strerror(-ret)); - } - } else { + if (ctx != nullptr) { LOG_INFO("allocating ctx: %lu", (uint64_t)ctx); - - ctx_map[thread_id] = ctx; } - + ctx_map[thread_id] = ctx; lk.unlock(); #endif } @@ -262,11 +415,8 @@ void LinuxAlignedFileReader::deregister_thread() { ctx_map.erase(it); } - // io_destroy is a syscall; keep it outside the lock to avoid blocking others - if (ailego::IOBackend::Instance().available() != - ailego::IOBackendType::kPread) { - LibAioLoader::Instance().io_destroy(ctx); - } + // Teardown is a syscall; keep it outside the lock to avoid blocking others. + destroy_io_ctx(ctx); LOG_INFO("returned ctx from thread"); #endif } @@ -274,13 +424,8 @@ void LinuxAlignedFileReader::deregister_thread() { void LinuxAlignedFileReader::deregister_all_threads() { #if (defined(__linux) || defined(__linux__)) std::unique_lock lk(ctx_mut); - bool aio_available = ailego::IOBackend::Instance().available() != - ailego::IOBackendType::kPread; for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { - IOContext ctx = x->second; - if (aio_available) { - LibAioLoader::Instance().io_destroy(ctx); - } + destroy_io_ctx(x->second); } ctx_map.clear(); #endif diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index a1cb7c91a..7eaf41e8b 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -18,6 +18,7 @@ #include #if (defined(__linux) || defined(__linux__)) +#include // raw-syscall io_uring wrapper (IoUringRing) #include // dlopen-based libaio wrapper #endif @@ -31,7 +32,31 @@ namespace zvec { namespace core { #if (defined(__linux) || defined(__linux__)) -typedef io_context_t IOContext; + +// IoBackend holds the per-thread I/O context for whichever async backend +// was successfully initialised at setup time. The priority is: +// 1. io_uring (raw kernel syscalls — zero dependency) +// 2. libaio (dlopen — soft dependency) +// 3. pread (always available — synchronous fallback) +// +// IOContext is a *pointer* to IoBackend, which preserves the existing +// sentinel conventions: nullptr means uninitialised and (IOContext)-1 is +// the invalid-handle sentinel returned by get_ctx() for unregistered +// threads. +struct IoBackend { + enum Backend : uint8_t { + NONE = 0, // synchronous pread + IO_URING = 1, // io_uring via raw syscalls + LIBAIO = 2, // libaio via dlopen + }; + + Backend backend{NONE}; + IoUringRing ring{}; + io_context_t aio_ctx{nullptr}; +}; + +typedef IoBackend *IOContext; + #else typedef uint32_t IOContext; #endif diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 5af446733..9f06b12f5 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -15,6 +15,7 @@ #include #include #include +#include #include #include "algorithm/diskann/diskann_params.h" #include "holder_builder.h" @@ -271,4 +272,16 @@ int DiskAnnIndex::Merge(const std::vector &indexes, return 0; } +ailego::IOBackendType DiskAnnIndex::io_backend_type() const { + auto &backend = ailego::IOBackend::Instance(); + ailego::IOBackendType type = backend.type(); + if (type == ailego::IOBackendType::kPread) { + LOG_WARN( + "Only synchronous pread() is available. Install libaio " + "(e.g. 'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " + "for async I/O support — performance will be degraded without it."); + } + return type; +} + } // namespace zvec::core_interface \ No newline at end of file diff --git a/src/include/zvec/ailego/io/io_backend.h b/src/include/zvec/ailego/io/io_backend.h index 008b01514..410452f4f 100644 --- a/src/include/zvec/ailego/io/io_backend.h +++ b/src/include/zvec/ailego/io/io_backend.h @@ -12,13 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. -// I/O backend type enum. +// I/O backend abstraction — public, dependency-free header. // -// This is the public, dependency-free part of the I/O backend abstraction. -// It defines the IOBackendType enum and the convenience helpers -// current_io_backend_type() / current_io_backend_description() so that -// public headers can reference IOBackendType without pulling in the -// internal IOBackend singleton or libaio_loader. +// This is the single io_backend.h in the project. It defines the +// IOBackendType enum, the IOBackendTypeName() helper, and the convenience +// helpers current_io_backend_type() / current_io_backend_description() so that +// public headers can reference IOBackendType without pulling in the internal +// IOBackend singleton, libaio_loader, or the io_uring kernel ABI headers. +// +// The IOBackend singleton (which probes io_uring / libaio at runtime) lives in +// the internal header ailego/io/io_backend_def.h. #pragma once @@ -29,13 +32,17 @@ namespace zvec { namespace ailego { // Supported I/O backend types. +// +// Numeric values are part of the C ABI (see zvec_io_backend_type_t in c_api.h): +// kPread = 0, kLibAio = 1, kIoUring = 2. enum class IOBackendType { - kPread, // Synchronous pread() — no async I/O - kLibAio, // libaio loaded at runtime via dlopen() + kPread = 0, // Synchronous pread() — no async I/O + kLibAio = 1, // libaio loaded at runtime via dlopen() + kIoUring = 2, // io_uring via raw kernel syscalls (zero dependency) }; // Returns the currently active I/O backend type. -// Triggers backend initialization on first call (libaio > pread). +// Triggers backend initialization on first call (io_uring > libaio > pread). IOBackendType current_io_backend_type(); // Returns a human-readable description of the currently active I/O backend. diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index 2d048689b..7e18886b4 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -789,6 +789,8 @@ typedef uint32_t zvec_io_backend_type_t; 0 /**< Synchronous pread() \u2014 no async I/O */ #define ZVEC_IO_BACKEND_TYPE_LIBAIO \ 1 /**< libaio loaded at runtime via dlopen() */ +#define ZVEC_IO_BACKEND_TYPE_IO_URING \ + 2 /**< io_uring via raw kernel syscalls (zero dependency) */ /** * @brief Get the current I/O backend type for DiskAnn async disk reads. @@ -796,7 +798,8 @@ typedef uint32_t zvec_io_backend_type_t; * Pure introspection \u2014 no side effects, no install hints. * * @return zvec_io_backend_type_t The loaded backend type - * (ZVEC_IO_BACKEND_TYPE_LIBAIO or ZVEC_IO_BACKEND_TYPE_PREAD). + * (ZVEC_IO_BACKEND_TYPE_IO_URING, ZVEC_IO_BACKEND_TYPE_LIBAIO, + * or ZVEC_IO_BACKEND_TYPE_PREAD). */ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void); @@ -805,7 +808,7 @@ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void); * * @param type The backend type code. * @return Thread-local string valid until the next call on this thread; - * "libaio", "pread", or "unknown". + * "io_uring", "libaio", "pread", or "unknown". */ ZVEC_EXPORT const char *ZVEC_CALL zvec_get_io_backend_type_name(zvec_io_backend_type_t type);