From c571f37a104c6430797de10bf10692b105a4d819 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Wed, 25 Feb 2026 14:59:48 +0100 Subject: [PATCH 01/10] feat(python): lay down the python structure for pyhmll package --- lib/python/loader.cpp | 48 +++++++++++++++++++++ lib/python/loader.hpp | 6 ++- lib/python/pyhmll/__init__.py | 62 +++++++++++++++++++++++++++- lib/python/pyhmll/torch.py | 78 +++++++++++++++++++++++++++++++++++ lib/python/pyproject.toml | 2 +- lib/python/safetensors.cpp | 30 +++++++++++++- 6 files changed, 220 insertions(+), 6 deletions(-) create mode 100644 lib/python/pyhmll/torch.py diff --git a/lib/python/loader.cpp b/lib/python/loader.cpp index 6ac912e..864f59e 100644 --- a/lib/python/loader.cpp +++ b/lib/python/loader.cpp @@ -95,6 +95,53 @@ size_t WeightLoader::fetch(const int iofile, const size_t offset, const uintptr_ } } +size_t WeightLoader::fetchv(const int iofile, const std::vector>& ranges, const uintptr_t dst) const +{ + nb::gil_scoped_release release; + + const auto ctx = ctx_.get(); + const auto dev = device(); + const auto* fetcher = ctx->fetcher; + + if (ranges.empty()) + return 0; + + ssize_t res = 0; + if (fetcher->fetchv_range_impl_) { + std::vector dsts(ranges.size()); + std::vector offsets(ranges.size()); + size_t dst_offset = 0; + for (size_t i = 0; i < ranges.size(); ++i) { + const auto [start, end] = ranges[i]; + const size_t nbytes = end - start; + dsts[i] = {nbytes, reinterpret_cast(dst + dst_offset), dev}; + offsets[i] = start; + dst_offset += nbytes; + } + res = hmll_fetchv(ctx, iofile, dsts.data(), offsets.data(), ranges.size()); + if (res < 0) { + const std::string err = hmll_strerr(ctx_->error); + throw std::runtime_error(fmt::format(PYHMLL_ERR_FETCH, err)); + } + return static_cast(res); + } + + /* Fallback: sequential fetch when fetchv is not implemented (e.g. Win32 mmap) */ + size_t total = 0; + size_t dst_offset = 0; + for (const auto& [start, end] : ranges) { + const size_t nbytes = end - start; + const hmll_iobuf_t buf = {nbytes, reinterpret_cast(dst + dst_offset), dev}; + if (res = hmll_fetch(ctx, iofile, &buf, start); res <= 0) { + const std::string err = hmll_strerr(ctx_->error); + throw std::runtime_error(fmt::format(PYHMLL_ERR_FETCH, err)); + } + total += static_cast(res); + dst_offset += nbytes; + } + return total; +} + void init_loader(nb::module_& m) { nb::class_(m, "Device", R"pbdoc(Define all the targetable devices)pbdoc") @@ -149,6 +196,7 @@ void init_loader(nb::module_& m) .def_prop_ro("kind", &WeightLoader::kind) .def("afetch", &WeightLoader::afetch) .def("fetch", &WeightLoader::fetch) + .def("fetchv", &WeightLoader::fetchv, "iofile"_a.sig("int"), "ranges"_a.sig("list[tuple[int, int]]"), "dst"_a.sig("int")) .def("__repr__", [](const WeightLoader& self) { return fmt::format(FMT_COMPILE("WeightLoader(kind={}, device={}})"), self.kind(), self.device()); diff --git a/lib/python/loader.hpp b/lib/python/loader.hpp index 4e2e850..23bbaa0 100644 --- a/lib/python/loader.hpp +++ b/lib/python/loader.hpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include "hmll/hmll.h" @@ -43,8 +44,9 @@ class WeightLoader [[nodiscard]] size_t fetch(int iofile, size_t offset, uintptr_t dst, size_t size) const; - [[nodiscard]] nb::ndarray, nb::c_contig> - fetchv(int iofile, const std::vector>& ranges, hmll_dtype_t dtype) const; + /** Batched fetch: write multiple byte ranges from file into a single pre-allocated buffer. */ + [[nodiscard]] size_t + fetchv(int iofile, const std::vector>& ranges, uintptr_t dst) const; }; #endif // PYHMLL_FETCHER_HPP diff --git a/lib/python/pyhmll/__init__.py b/lib/python/pyhmll/__init__.py index 4e4640b..0571cae 100644 --- a/lib/python/pyhmll/__init__.py +++ b/lib/python/pyhmll/__init__.py @@ -1 +1,61 @@ -from .pyhmll_impl import * \ No newline at end of file +""" +Python layer on top of _pyhmll_impl. + +Provides device-agnostic API using torch.device and dtype conversion to torch.dtype. +""" +from __future__ import annotations + +from pathlib import Path +from typing import Any + +try: + from _pyhmll_impl import ( + Backend, + Device, + SafetensorsAccessor, + dtype, + safetensors as _safetensors_impl, + ) +except ImportError as e: + raise ImportError( + "pyhmll requires the _pyhmll_impl C extension. " + "Build the hmll project with the Python bindings enabled." + ) from e + +from pyhmll.torch import as_dtype + + +def safetensors( + path: str | Path, + device: Device = None, + is_sharded: bool = False, + backend: Any = None, +) -> SafetensorsAccessor: + """ + Open a safetensors file or sharded index for fast loading. + + Args: + path: Path to a single .safetensors file or to model.safetensors.index.json + for sharded checkpoints. + device: Target device. If None, defaults to "cpu". + is_sharded: True if path is the path to the index file of a sharded checkpoint. + backend: I/O backend (e.g. Backend.AUTO, Backend.IO_URING, Backend.MMAP). + Pass None for default. + + Returns: + SafetensorsAccessor context manager for reading tensors. + """ + device_ = device or Device.CPU + backend_ = backend if backend is not None else Backend.IO_URING + path_ = str(path) + return _safetensors_impl(path_, device_, is_sharded, backend_) + + +__all__ = [ + "as_dtype", + "Backend", + "Device", + "SafetensorsAccessor", + "dtype", + "safetensors", +] diff --git a/lib/python/pyhmll/torch.py b/lib/python/pyhmll/torch.py new file mode 100644 index 0000000..43271ea --- /dev/null +++ b/lib/python/pyhmll/torch.py @@ -0,0 +1,78 @@ +""" +PyTorch-specific utilities for pyhmll. +""" +from __future__ import annotations + +import torch + +# Import the low-level dtype enum from the C extension +try: + from _pyhmll_impl import dtype as _hmll_dtype + from _pyhmll_impl import Device +except ImportError: + _hmll_dtype = None + + +def as_dtype(hmll_dtype): + """ + Convert an hmll dtype enum value to torch.dtype. + + Args: + hmll_dtype: Value from _pyhmll_impl.dtype (e.g. dtype.BFLOAT16). + + Returns: + Corresponding torch.dtype. + """ + if _hmll_dtype is None: + raise ImportError("pyhmll C extension (_pyhmll_impl) not available") + + match hmll_dtype: + case _hmll_dtype.BOOL: + return torch.bool + case _hmll_dtype.BFLOAT16: + return torch.bfloat16 + case _hmll_dtype.COMPLEX: + return torch.complex64 + case _hmll_dtype.FLOAT16: + return torch.float16 + case _hmll_dtype.FLOAT32: + return torch.float32 + case _hmll_dtype.FLOAT64: + return torch.float64 + case _hmll_dtype.FLOAT8_E8M0: + return getattr(torch, "float8_e8m0fn", torch.uint8) + case _hmll_dtype.FLOAT8_E4M3: + return getattr(torch, "float8_e4m3fn", torch.float32) + case _hmll_dtype.FLOAT8_E5M2: + return getattr(torch, "float8_e5m2", torch.float32) + case _hmll_dtype.SIGNED_INT8: + return torch.int8 + case _hmll_dtype.SIGNED_INT16: + return torch.int16 + case _hmll_dtype.SIGNED_INT32: + return torch.int32 + case _hmll_dtype.SIGNED_INT64: + return torch.int64 + case _hmll_dtype.UNSIGNED_INT8: + return torch.uint8 + case _hmll_dtype.UNSIGNED_INT16: + return torch.uint16 + case _hmll_dtype.UNSIGNED_INT32: + return torch.uint32 + case _hmll_dtype.UNSIGNED_INT64: + return torch.uint64 + case _hmll_dtype.UNKNOWN: + raise ValueError(f"No torch.dtype mapping for hmll dtype {hmll_dtype}") + + +def device_to_hmll(device: torch.device) -> Device: + """ + Convert a torch.device to a hmll Device enum value. + """ + match device.type: + case "cuda": + return Device.CUDA + case "cpu": + return Device.CPU + case _: + raise ValueError(f"Unsupported device for pyhmll: {device!r}") \ No newline at end of file diff --git a/lib/python/pyproject.toml b/lib/python/pyproject.toml index a5b19d2..1c09418 100644 --- a/lib/python/pyproject.toml +++ b/lib/python/pyproject.toml @@ -3,7 +3,7 @@ name = "pyhmll" version = "0.1.0" description = "Hugging Face Models Loading Library" readme = "../../README.md" -requires-python = ">=3.9" +requires-python = ">=3.10" authors = [ { name = "Morgan Funtowicz", email = "morgan@hf.co" }, ] diff --git a/lib/python/safetensors.cpp b/lib/python/safetensors.cpp index cc8e6d3..8a1cbd5 100644 --- a/lib/python/safetensors.cpp +++ b/lib/python/safetensors.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include "loader.hpp" @@ -155,6 +156,31 @@ class SafetensorsAccessor // Delegate to WeightLoader for actual fetching return loader_->fetch(iofile, specs.start, dst, size); } + + /** Fetch only the given element ranges into dst. Dtype is taken from registry. */ + [[nodiscard]] size_t fetchv(const std::string& name, const std::vector>& ranges, const uintptr_t dst) const + { + const auto registry = registry_.get(); + const hmll_lookup_result lookup = hmll_lookup_tensor(loader_->context(), registry, name.c_str()); + if (lookup.specs == nullptr) + throw nb::key_error(name.c_str()); + + const hmll_tensor_specs_t* specs = lookup.specs; + const int iofile = lookup.file; + const size_t numel = hmll_numel(specs); + const size_t nbytes = specs->end - specs->start; + const size_t nbits = (numel > 0) ? (nbytes / numel) : 0; + + std::vector> byte_ranges; + byte_ranges.reserve(ranges.size()); + for (const auto& [start, end] : ranges) { + const size_t file_start = specs->start + start * nbits; + const size_t file_end = specs->start + end * nbits; + byte_ranges.emplace_back(file_start, file_end); + } + + return loader_->fetchv(iofile, byte_ranges, dst); + } }; void init_safetensors(nb::module_& m) @@ -180,8 +206,8 @@ void init_safetensors(nb::module_& m) .def("values", &SafetensorsAccessor::specs, nb::rv_policy::reference_internal) .def("items", &SafetensorsAccessor::named_specs, nb::rv_policy::reference_internal) .def("afetch", &SafetensorsAccessor::afetch) - .def("fetch", &SafetensorsAccessor::fetch); - + .def("fetch", &SafetensorsAccessor::fetch) + .def("fetchv", &SafetensorsAccessor::fetchv); m.def("safetensors", [](const std::filesystem::path& path, const hmll_device_t device, const bool is_sharded, const hmll_fetcher_kind_t backend) { return new SafetensorsAccessor(path, device, is_sharded, backend); }, nb::rv_policy::take_ownership, "path"_a, "device"_a, "is_sharded"_a = false, "backend"_a = HMLL_FETCHER_AUTO); From be346bca18dd06f5b8bfc913213d597af3d73aa8 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Wed, 25 Feb 2026 15:01:33 +0100 Subject: [PATCH 02/10] make fetcher use the cca window value and not the static batch size --- CMakeLists.txt | 6 +++--- include/hmll/linux/backend/iouring.h | 4 ---- lib/linux/backend/iouring.c | 8 ++++---- 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index f32b82c..fc85808 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -142,9 +142,9 @@ if(${HMLL_ENABLE_SAFETENSORS}) add_compile_definitions(__HMLL_SAFETENSORS_ENABLED__=1) add_compile_definitions(__HMLL_TENSORS_ENABLED__=1) fetchcontent_declare( - yyjson - GIT_REPOSITORY https://github.com/ibireme/yyjson.git - GIT_TAG master + yyjson + GIT_REPOSITORY https://github.com/ibireme/yyjson.git + GIT_TAG master ) fetchcontent_makeavailable(yyjson) diff --git a/include/hmll/linux/backend/iouring.h b/include/hmll/linux/backend/iouring.h index 85c6a8d..49eff45 100644 --- a/include/hmll/linux/backend/iouring.h +++ b/include/hmll/linux/backend/iouring.h @@ -11,10 +11,6 @@ #define HMLL_URING_BUFFER_SIZE (512U * 1024) #endif -#ifndef HMLL_URING_CQE_BATCH_SIZE -#define HMLL_URING_CQE_BATCH_SIZE 16 -#endif - #include #include "hmll/types.h" diff --git a/lib/linux/backend/iouring.c b/lib/linux/backend/iouring.c index 661d7e0..0f83cf5 100644 --- a/lib/linux/backend/iouring.c +++ b/lib/linux/backend/iouring.c @@ -203,7 +203,7 @@ static ssize_t hmll_io_uring_fetch_range_impl( size_t n_dma = 0; size_t b_read = 0; size_t b_submitted = 0; - struct io_uring_cqe *cqes[HMLL_URING_CQE_BATCH_SIZE]; + struct io_uring_cqe *cqes[HMLL_URING_QUEUE_DEPTH]; struct io_uring_sqe *sqe = NULL; int slot; @@ -249,7 +249,7 @@ static ssize_t hmll_io_uring_fetch_range_impl( } unsigned count = 0; - while ((count = io_uring_peek_batch_cqe(&fetcher->ioring, cqes, HMLL_URING_CQE_BATCH_SIZE)) > 0) { + while ((count = io_uring_peek_batch_cqe(&fetcher->ioring, cqes, fetcher->iocca.window)) > 0) { for (unsigned i = 0; i < count; i++) { const struct io_uring_cqe *cqe = cqes[i]; @@ -340,7 +340,7 @@ static ssize_t hmll_io_uring_fetchv_range_impl( const unsigned char is_cuda = hmll_device_is_cuda(dsts[0].device); size_t n_in_flight = 0, nbytes = 0, active_cursor = 0; - struct io_uring_cqe *cqes[HMLL_URING_CQE_BATCH_SIZE]; + struct io_uring_cqe *cqes[HMLL_URING_QUEUE_DEPTH]; while (n_active > 0 || n_in_flight > 0) { while (n_active > 0) { @@ -416,7 +416,7 @@ static ssize_t hmll_io_uring_fetchv_range_impl( if (nwait > 0) hmll_io_uring_cca_update(&fetcher->iocca, HMLL_URING_BUFFER_SIZE * nwait, ts_start, ts_end); unsigned count; - while ((count = io_uring_peek_batch_cqe(&fetcher->ioring, cqes, HMLL_URING_CQE_BATCH_SIZE)) > 0) { + while ((count = io_uring_peek_batch_cqe(&fetcher->ioring, cqes, fetcher->iocca.window)) > 0) { for (unsigned i = 0; i < count; i++) { const struct io_uring_cqe *cqe = cqes[i]; const uint64_t data = cqe->user_data; From b1a28446b00d6d569fbde3d00e66268ebd8bf5f3 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Thu, 26 Feb 2026 14:40:32 +0100 Subject: [PATCH 03/10] more wip integration --- lib/linux/backend/iouring.c | 207 ++++++++++++++++++++---------------- lib/python/safetensors.cpp | 6 +- 2 files changed, 116 insertions(+), 97 deletions(-) diff --git a/lib/linux/backend/iouring.c b/lib/linux/backend/iouring.c index 0f83cf5..58db964 100644 --- a/lib/linux/backend/iouring.c +++ b/lib/linux/backend/iouring.c @@ -1,35 +1,19 @@ #include -#include #include "hmll/hmll.h" -#include "hmll/cuda.h" #include "hmll/memory.h" #include "hmll/linux/backend/iouring.h" #include "sys/mman.h" -#include #define HMLL_IO_URING_ADVISORY_FLAG UINT64_MAX #if defined(__HMLL_CUDA_ENABLED__) +#include "hmll/cuda.h" #include #include #endif -static inline int hmll_io_uring_get_setup_flags(void) -{ - int flags = IORING_SETUP_SQPOLL; - - // retrieve the current kernel version so we can adjust io_uring flags - struct utsname unamedata; - uname(&unamedata); - - int major, minor, revision = 0; - if (sscanf(unamedata.release, "%d.%d.%d", &major, &minor, &revision)) { - if (major >= 6) flags |= IORING_SETUP_SINGLE_ISSUER; - } - - return flags; -} +static inline int hmll_io_uring_get_setup_flags(void) { return IORING_SETUP_SQPOLL; } static struct hmll_error hmll_io_uring_register_staging_buffers( struct hmll *ctx, @@ -190,7 +174,7 @@ static inline void hmll_io_uring_handle_completion( #endif } -static ssize_t hmll_io_uring_fetch_range_impl( +static ssize_t hmll_io_uring_fetch_impl( struct hmll *ctx, const int iofile, const struct hmll_iobuf *dst, @@ -208,7 +192,7 @@ static ssize_t hmll_io_uring_fetch_range_impl( struct io_uring_sqe *sqe = NULL; int slot; if ((sqe = io_uring_get_sqe(&fetcher->ioring))) { - io_uring_prep_fadvise(sqe, iofile, offset, dst->size, POSIX_FADV_SEQUENTIAL | POSIX_FADV_WILLNEED); + io_uring_prep_fadvise(sqe, iofile, offset, dst->size, POSIX_FADV_WILLNEED); io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); io_uring_sqe_set_data64(sqe, HMLL_IO_URING_ADVISORY_FLAG); } @@ -259,7 +243,7 @@ static ssize_t hmll_io_uring_fetch_range_impl( --n_dma; if (unlikely(cqe->res < 0)) { ctx->error = HMLL_SYS_ERR(-cqe->res); - io_uring_cq_advance(&fetcher->ioring, i + 1); + io_uring_cq_advance(&fetcher->ioring, count); return -1; } @@ -275,7 +259,7 @@ static ssize_t hmll_io_uring_fetch_range_impl( return (ssize_t)b_read; } -static ssize_t hmll_io_uring_fetchv_range_impl( +static ssize_t hmll_io_uring_fetchv_impl( struct hmll *ctx, const int iofile, const struct hmll_iobuf *dsts, @@ -283,6 +267,7 @@ static ssize_t hmll_io_uring_fetchv_range_impl( const size_t n ) { if (hmll_check(ctx->error)) return -1; + if (unlikely(n == 0)) return 0; struct hmll_io_uring *fetcher = ctx->fetcher->backend_impl_; @@ -298,21 +283,20 @@ static ssize_t hmll_io_uring_fetchv_range_impl( _Alignas(16) uint8_t stack_mem[8192]; - const size_t state_mem_req = sizeof(struct fetch_state) * n; - const size_t idx_mem_req = sizeof(uint32_t) * n; + /* + * Align each sub-array to the requirement of the next type to avoid UB + * on strict-alignment targets (e.g. slot_offsets needs 8-byte alignment). + */ + const size_t state_mem_req = (sizeof(struct fetch_state) * n + (_Alignof(uint32_t) - 1)) & ~(_Alignof(uint32_t) - 1); + const size_t idx_mem_req = (sizeof(uint32_t) * n + (_Alignof(size_t) - 1)) & ~(_Alignof(size_t) - 1); const size_t slot_mem_req = sizeof(size_t) * HMLL_URING_QUEUE_DEPTH; const size_t total_req = state_mem_req + idx_mem_req + slot_mem_req; if (likely(total_req <= sizeof(stack_mem))) { uint8_t *ptr = stack_mem; - - states = (struct fetch_state *)ptr; - ptr += state_mem_req; - - active_indices = (uint32_t *)ptr; - ptr += idx_mem_req; - - slot_offsets = (size_t *)ptr; + states = (struct fetch_state *)ptr; ptr += state_mem_req; + active_indices = (uint32_t *)ptr; ptr += idx_mem_req; + slot_offsets = (size_t *)ptr; } else { states = calloc(1, total_req); if (unlikely(!states)) { @@ -320,71 +304,71 @@ static ssize_t hmll_io_uring_fetchv_range_impl( return -1; } active_indices = (uint32_t *)((char *)states + state_mem_req); - slot_offsets = (size_t *)((char *)active_indices + idx_mem_req); + slot_offsets = (size_t *)((char *)active_indices + idx_mem_req); } size_t n_active = 0; for (size_t i = 0; i < n; ++i) { - states[i].submitted = 0; - states[i].size = dsts[i].size; - states[i].fadvise_sent = false; - - if (dsts[i].size > 0) { - active_indices[n_active++] = i; - } + states[i].submitted = 0; + states[i].size = dsts[i].size; + states[i].fadvise_sent = 0; + if (dsts[i].size > 0) + active_indices[n_active++] = (uint32_t)i; } - const uint64_t BIT_FADVISE = 1ULL << 63; - const uint64_t SHIFT_RANGE = 32; - const uint64_t MASK_SLOT = 0xFFFFFFFFULL; const unsigned char is_cuda = hmll_device_is_cuda(dsts[0].device); - size_t n_in_flight = 0, nbytes = 0, active_cursor = 0; struct io_uring_cqe *cqes[HMLL_URING_QUEUE_DEPTH]; while (n_active > 0 || n_in_flight > 0) { - while (n_active > 0) { - struct io_uring_sqe *sqe = io_uring_get_sqe(&fetcher->ioring); - if (!sqe) break; + /* Eagerly reclaim completed CUDA staging slots each outer iteration */ + hmll_io_uring_reclaim_slots(fetcher, dsts[0].device); + /* Submit as many read SQEs as possible, round-robining across active buffers */ + while (n_active > 0) { if (active_cursor >= n_active) active_cursor = 0; - const uint32_t current_idx = active_indices[active_cursor]; - struct fetch_state *st = &states[current_idx]; + const uint32_t bidx = active_indices[active_cursor]; + struct fetch_state *st = &states[bidx]; + /* Send fadvise hint once per buffer; active_cursor stays unchanged so + the same buffer gets its data read SQE on the very next iteration. */ if (unlikely(!st->fadvise_sent)) { - io_uring_prep_fadvise(sqe, iofile, offsets[current_idx], st->size, POSIX_FADV_SEQUENTIAL | POSIX_FADV_WILLNEED); - io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); - io_uring_sqe_set_data64(sqe, BIT_FADVISE); + struct io_uring_sqe *fadvise_sqe = io_uring_get_sqe(&fetcher->ioring); + if (!fadvise_sqe) break; // ring full, submit what we have and retry + io_uring_prep_fadvise(fadvise_sqe, iofile, offsets[bidx], st->size, POSIX_FADV_WILLNEED); + io_uring_sqe_set_flags(fadvise_sqe, IOSQE_FIXED_FILE); + io_uring_sqe_set_data64(fadvise_sqe, BIT_FADVISE); st->fadvise_sent = 1; continue; } - int slot = hmll_io_uring_slot_find_available(fetcher->iobusy); - if (slot == -1) { - hmll_io_uring_reclaim_slots(fetcher, dsts[0].device); - slot = hmll_io_uring_slot_find_available(fetcher->iobusy); - if (slot == -1) break; - } + const int slot = hmll_io_uring_slot_find_available(fetcher->iobusy); + if (slot == -1) break; - hmll_io_uring_slot_set_busy(&fetcher->iobusy, slot); + struct io_uring_sqe *sqe = io_uring_get_sqe(&fetcher->ioring); + if (!sqe) break; + hmll_io_uring_slot_set_busy(&fetcher->iobusy, slot); slot_offsets[slot] = st->submitted; - const size_t remaining = st->size - st->submitted; - const size_t to_read = remaining < HMLL_URING_BUFFER_SIZE ? remaining : HMLL_URING_BUFFER_SIZE; - const size_t file_offset = offsets[current_idx] + st->submitted; - - hmll_io_uring_prep_sqe( - fetcher, - dsts[current_idx].device, - sqe, - (char *)dsts[current_idx].ptr + st->submitted, - file_offset, - to_read, - iofile, - slot - ); - - io_uring_sqe_set_data64(sqe, ((uint64_t)current_idx << SHIFT_RANGE) | slot); + + const size_t remaining = st->size - st->submitted; + const size_t to_read = remaining < HMLL_URING_BUFFER_SIZE ? remaining : HMLL_URING_BUFFER_SIZE; + const size_t file_offset = offsets[bidx] + st->submitted; + + /* CPU: read directly into the destination buffer. + CUDA: read into a pinned staging buffer; the async memcpy to GPU + is dispatched on CQE completion below. */ +#if defined(__HMLL_CUDA_ENABLED__) + if (is_cuda) + io_uring_prep_read_fixed(sqe, iofile, fetcher->iovecs[slot].iov_base, to_read, file_offset, slot); + else + io_uring_prep_read(sqe, iofile, (char *)dsts[bidx].ptr + st->submitted, to_read, file_offset); +#else + io_uring_prep_read(sqe, iofile, (char *)dsts[bidx].ptr + st->submitted, to_read, file_offset); +#endif + io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); + /* Pack buffer index and staging slot into user_data for CQE dispatch */ + io_uring_sqe_set_data64(sqe, ((uint64_t)bidx << SHIFT_RANGE) | (uint64_t)slot); st->submitted += to_read; n_in_flight++; @@ -397,13 +381,43 @@ static ssize_t hmll_io_uring_fetchv_range_impl( } } - size_t nwait = 0; - if (n_in_flight > 0) { - nwait = (n_in_flight < fetcher->iocca.window) ? n_in_flight : fetcher->iocca.window; - } else if (n_active == 0) { - break; + if (n_in_flight == 0) { + if (n_active == 0) break; + + /* + * n_in_flight == 0 but n_active > 0: the inner loop couldn't claim + * any staging slot. For CUDA this means all 128 slots are occupied + * by pending cudaMemcpyAsync operations whose events haven't fired + * yet. Spinning here (nwait=0 path) would burn 100 % CPU for + * hundreds of milliseconds on large models. + * + * Fix: flush any queued fadvise SQEs, then block on the first + * outstanding CUDA event so we reclaim at least one slot before + * retrying. For CPU this case is transient (SQPOLL catching up), so + * the io_uring_submit() wake-up is sufficient. + */ + io_uring_submit(&fetcher->ioring); + +#if defined(__HMLL_CUDA_ENABLED__) + if (is_cuda && hmll_io_uring_slot_find_available(fetcher->iobusy) == -1) { + struct hmll_io_uring_cuda_context *dctx = fetcher->device_ctx; + for (size_t i = 0; i < HMLL_URING_QUEUE_DEPTH; ++i) { + struct hmll_io_uring_cuda_context *cd = dctx + i; + if (hmll_io_uring_slot_is_busy(fetcher->iobusy, i) && + cd->state == HMLL_CUDA_STREAM_MEMCPY) { + cudaEventSynchronize(cd->done); + hmll_io_uring_cuda_stream_set_idle(&cd->state); + hmll_io_uring_slot_set_available(&fetcher->iobusy, i); + break; + } + } + } +#endif + continue; } + const size_t nwait = n_in_flight < fetcher->iocca.window ? n_in_flight : fetcher->iocca.window; + struct timespec ts_start, ts_end; clock_gettime(CLOCK_MONOTONIC, &ts_start); @@ -413,7 +427,7 @@ static ssize_t hmll_io_uring_fetchv_range_impl( } clock_gettime(CLOCK_MONOTONIC, &ts_end); - if (nwait > 0) hmll_io_uring_cca_update(&fetcher->iocca, HMLL_URING_BUFFER_SIZE * nwait, ts_start, ts_end); + hmll_io_uring_cca_update(&fetcher->iocca, HMLL_URING_BUFFER_SIZE * nwait, ts_start, ts_end); unsigned count; while ((count = io_uring_peek_batch_cqe(&fetcher->ioring, cqes, fetcher->iocca.window)) > 0) { @@ -421,45 +435,50 @@ static ssize_t hmll_io_uring_fetchv_range_impl( const struct io_uring_cqe *cqe = cqes[i]; const uint64_t data = cqe->user_data; - if (unlikely(data & BIT_FADVISE)) continue; + /* Skip fadvise completions; use exact equality to avoid false + positives if future encodings ever set high bits. */ + if (unlikely(data == BIT_FADVISE)) continue; n_in_flight--; if (unlikely(cqe->res < 0)) { - ctx->error = HMLL_ERR(HMLL_ERR_IO_ERROR); + ctx->error = HMLL_SYS_ERR(-cqe->res); io_uring_cq_advance(&fetcher->ioring, count); goto cleanup; } nbytes += cqe->res; - const uint32_t s_idx = (uint32_t)(data & MASK_SLOT); + const uint32_t slot = (uint32_t)(data & MASK_SLOT); + const uint32_t bidx = (uint32_t)(data >> SHIFT_RANGE); if (!is_cuda) { - hmll_io_uring_slot_set_available(&fetcher->iobusy, s_idx); + hmll_io_uring_slot_set_available(&fetcher->iobusy, slot); } #if defined(__HMLL_CUDA_ENABLED__) else { - const uint32_t r_idx = (uint32_t)(data >> SHIFT_RANGE); - struct hmll_io_uring_cuda_context *cctx = &((struct hmll_io_uring_cuda_context *)fetcher->device_ctx)[s_idx]; - void *to = (char *)dsts[r_idx].ptr + slot_offsets[s_idx]; - void *from = fetcher->iovecs[s_idx].iov_base; - + struct hmll_io_uring_cuda_context *cctx = + &((struct hmll_io_uring_cuda_context *)fetcher->device_ctx)[slot]; + void *from = fetcher->iovecs[slot].iov_base; + void *to = (char *)dsts[bidx].ptr + slot_offsets[slot]; cudaMemcpyAsync(to, from, cqe->res, cudaMemcpyHostToDevice, cctx->stream); cudaEventRecord(cctx->done, cctx->stream); hmll_io_uring_cuda_stream_set_memcpy(&cctx->state); } +#else + (void)bidx; #endif } io_uring_cq_advance(&fetcher->ioring, count); } } - if ((unsigned char*)states != stack_mem) free(states); + hmll_io_uring_sync(dsts[0].device, fetcher); + if ((unsigned char *)states != stack_mem) free(states); return (ssize_t)nbytes; cleanup: - if ((unsigned char*)states != stack_mem) free(states); + if ((unsigned char *)states != stack_mem) free(states); return -1; } @@ -532,8 +551,8 @@ struct hmll_error hmll_io_uring_init(struct hmll *ctx, const struct hmll_device ctx->fetcher->kind = HMLL_FETCHER_IO_URING; ctx->fetcher->device = device; ctx->fetcher->backend_impl_ = backend; - ctx->fetcher->fetch_range_impl_ = hmll_io_uring_fetch_range_impl; - ctx->fetcher->fetchv_range_impl_ = hmll_io_uring_fetchv_range_impl; + ctx->fetcher->fetch_range_impl_ = hmll_io_uring_fetch_impl; + ctx->fetcher->fetchv_range_impl_ = hmll_io_uring_fetchv_impl; ctx->fetcher->backend_free = hmll_io_uring_destroy; } diff --git a/lib/python/safetensors.cpp b/lib/python/safetensors.cpp index 8a1cbd5..aa8ceab 100644 --- a/lib/python/safetensors.cpp +++ b/lib/python/safetensors.cpp @@ -69,7 +69,7 @@ class SafetensorsAccessor if (hmll_check(hmll_source_open(path_str.c_str(), &source))) throw std::runtime_error("Failed to open file: " + path_str); - const auto registry = std::make_shared(); + auto registry = std::make_shared(); auto ctx = std::make_unique(); if (const auto n_tensors = hmll_safetensors_populate_registry(ctx.get(), registry.get(), source, 0, 0); n_tensors == 0) { @@ -77,7 +77,7 @@ class SafetensorsAccessor throw std::runtime_error(fmt::format(FMT_COMPILE("Failed to read tensor definition in file {}: {}"), path, hmll_strerr(ctx->error))); } - auto sources = std::vector{source}; + auto sources = std::vector{source}; loader_ = std::make_unique(std::move(sources), device, std::move(ctx), backend); registry_ = std::move(registry); } @@ -157,7 +157,7 @@ class SafetensorsAccessor return loader_->fetch(iofile, specs.start, dst, size); } - /** Fetch only the given element ranges into dst. Dtype is taken from registry. */ + /** Fetch only the given element ranges into dst. Dtype is taken from the registry. */ [[nodiscard]] size_t fetchv(const std::string& name, const std::vector>& ranges, const uintptr_t dst) const { const auto registry = registry_.get(); From 7fb2cb38c43438aafb62bd85a146317ac5d21f6f Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 2 Mar 2026 14:54:54 +0100 Subject: [PATCH 04/10] cleanup of fetchv iouring.c --- lib/context.c | 1 + lib/linux/backend/iouring.c | 164 ++++++++++++++++-------------------- 2 files changed, 72 insertions(+), 93 deletions(-) diff --git a/lib/context.c b/lib/context.c index 6ceb63e..0ffc941 100644 --- a/lib/context.c +++ b/lib/context.c @@ -19,5 +19,6 @@ void hmll_destroy(struct hmll *ctx) if (ctx->fetcher) { ctx->fetcher->backend_free(ctx->fetcher->backend_impl_); free(ctx->fetcher); + ctx->fetcher = NULL; } } diff --git a/lib/linux/backend/iouring.c b/lib/linux/backend/iouring.c index 58db964..fbf72b8 100644 --- a/lib/linux/backend/iouring.c +++ b/lib/linux/backend/iouring.c @@ -270,48 +270,54 @@ static ssize_t hmll_io_uring_fetchv_impl( if (unlikely(n == 0)) return 0; struct hmll_io_uring *fetcher = ctx->fetcher->backend_impl_; + const int is_cuda = (dsts[0].device == HMLL_DEVICE_CUDA); - struct fetch_state { + /* user_data encoding for CQEs: high bit = fadvise (skip), else (bidx << 8) | slot */ + static const uint64_t FETCHV_FADVISE_TAG = 1ULL << 63; + static const unsigned FETCHV_BIDX_SHIFT = 8; + static const uint64_t FETCHV_SLOT_MASK = HMLL_URING_QUEUE_DEPTH - 1; + + struct fetchv_buf_state { size_t submitted; size_t size; unsigned char fadvise_sent; }; - struct fetch_state *states; + /* Scratch layout: [buf_states][active_indices][slot_offsets] */ + const size_t sz_state = (sizeof(struct fetchv_buf_state) * n + _Alignof(uint32_t) - 1) & ~(_Alignof(uint32_t) - 1); + const size_t sz_idx = (sizeof(uint32_t) * n + _Alignof(size_t) - 1) & ~(_Alignof(size_t) - 1); + const size_t sz_slot = sizeof(size_t) * HMLL_URING_QUEUE_DEPTH; + const size_t scratch_size = sz_state + sz_idx + sz_slot; + + _Alignas(16) uint8_t stack_scratch[8192]; + struct fetchv_buf_state *buf_states; uint32_t *active_indices; size_t *slot_offsets; + void *scratch_to_free = NULL; - _Alignas(16) uint8_t stack_mem[8192]; - - /* - * Align each sub-array to the requirement of the next type to avoid UB - * on strict-alignment targets (e.g. slot_offsets needs 8-byte alignment). - */ - const size_t state_mem_req = (sizeof(struct fetch_state) * n + (_Alignof(uint32_t) - 1)) & ~(_Alignof(uint32_t) - 1); - const size_t idx_mem_req = (sizeof(uint32_t) * n + (_Alignof(size_t) - 1)) & ~(_Alignof(size_t) - 1); - const size_t slot_mem_req = sizeof(size_t) * HMLL_URING_QUEUE_DEPTH; - const size_t total_req = state_mem_req + idx_mem_req + slot_mem_req; - - if (likely(total_req <= sizeof(stack_mem))) { - uint8_t *ptr = stack_mem; - states = (struct fetch_state *)ptr; ptr += state_mem_req; - active_indices = (uint32_t *)ptr; ptr += idx_mem_req; - slot_offsets = (size_t *)ptr; + if (scratch_size <= sizeof(stack_scratch)) { + uint8_t *p = stack_scratch; + buf_states = (struct fetchv_buf_state *)p; p += sz_state; + active_indices = (uint32_t *)p; p += sz_idx; + slot_offsets = (size_t *)p; } else { - states = calloc(1, total_req); - if (unlikely(!states)) { + void *p = calloc(1, scratch_size); + if (!p) { ctx->error = HMLL_ERR(HMLL_ERR_ALLOCATION_FAILED); return -1; } - active_indices = (uint32_t *)((char *)states + state_mem_req); - slot_offsets = (size_t *)((char *)active_indices + idx_mem_req); + scratch_to_free = p; + buf_states = (struct fetchv_buf_state *)p; p = (char *)p + sz_state; + active_indices = (uint32_t *)p; p = (char *)p + sz_idx; + slot_offsets = (size_t *)p; } + /* Build list of buffers that have bytes to read */ size_t n_active = 0; - for (size_t i = 0; i < n; ++i) { - states[i].submitted = 0; - states[i].size = dsts[i].size; - states[i].fadvise_sent = 0; + for (size_t i = 0; i < n; i++) { + buf_states[i].submitted = 0; + buf_states[i].size = dsts[i].size; + buf_states[i].fadvise_sent = 0; if (dsts[i].size > 0) active_indices[n_active++] = (uint32_t)i; } @@ -321,54 +327,47 @@ static ssize_t hmll_io_uring_fetchv_impl( struct io_uring_cqe *cqes[HMLL_URING_QUEUE_DEPTH]; while (n_active > 0 || n_in_flight > 0) { - /* Eagerly reclaim completed CUDA staging slots each outer iteration */ hmll_io_uring_reclaim_slots(fetcher, dsts[0].device); - /* Submit as many read SQEs as possible, round-robining across active buffers */ + /* Submit: round-robin over active buffers, send fadvise then chunked reads */ while (n_active > 0) { if (active_cursor >= n_active) active_cursor = 0; const uint32_t bidx = active_indices[active_cursor]; - struct fetch_state *st = &states[bidx]; - - /* Send fadvise hint once per buffer; active_cursor stays unchanged so - the same buffer gets its data read SQE on the very next iteration. */ - if (unlikely(!st->fadvise_sent)) { - struct io_uring_sqe *fadvise_sqe = io_uring_get_sqe(&fetcher->ioring); - if (!fadvise_sqe) break; // ring full, submit what we have and retry - io_uring_prep_fadvise(fadvise_sqe, iofile, offsets[bidx], st->size, POSIX_FADV_WILLNEED); - io_uring_sqe_set_flags(fadvise_sqe, IOSQE_FIXED_FILE); - io_uring_sqe_set_data64(fadvise_sqe, BIT_FADVISE); + struct fetchv_buf_state *st = &buf_states[bidx]; + + if (!st->fadvise_sent) { + struct io_uring_sqe *sqe = io_uring_get_sqe(&fetcher->ioring); + if (!sqe) break; + io_uring_prep_fadvise(sqe, iofile, offsets[bidx], st->size, POSIX_FADV_WILLNEED); + io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); + io_uring_sqe_set_data64(sqe, FETCHV_FADVISE_TAG); st->fadvise_sent = 1; continue; } const int slot = hmll_io_uring_slot_find_available(fetcher->iobusy); - if (slot == -1) break; - + if (slot < 0) break; struct io_uring_sqe *sqe = io_uring_get_sqe(&fetcher->ioring); if (!sqe) break; hmll_io_uring_slot_set_busy(&fetcher->iobusy, slot); slot_offsets[slot] = st->submitted; - const size_t remaining = st->size - st->submitted; - const size_t to_read = remaining < HMLL_URING_BUFFER_SIZE ? remaining : HMLL_URING_BUFFER_SIZE; - const size_t file_offset = offsets[bidx] + st->submitted; + const size_t remaining = st->size - st->submitted; + const size_t to_read = remaining < HMLL_URING_BUFFER_SIZE ? remaining : HMLL_URING_BUFFER_SIZE; + const size_t file_off = offsets[bidx] + st->submitted; + void *read_dst = (char *)dsts[bidx].ptr + st->submitted; - /* CPU: read directly into the destination buffer. - CUDA: read into a pinned staging buffer; the async memcpy to GPU - is dispatched on CQE completion below. */ #if defined(__HMLL_CUDA_ENABLED__) if (is_cuda) - io_uring_prep_read_fixed(sqe, iofile, fetcher->iovecs[slot].iov_base, to_read, file_offset, slot); + io_uring_prep_read_fixed(sqe, iofile, fetcher->iovecs[slot].iov_base, to_read, file_off, slot); else - io_uring_prep_read(sqe, iofile, (char *)dsts[bidx].ptr + st->submitted, to_read, file_offset); + io_uring_prep_read(sqe, iofile, read_dst, to_read, file_off); #else - io_uring_prep_read(sqe, iofile, (char *)dsts[bidx].ptr + st->submitted, to_read, file_offset); + io_uring_prep_read(sqe, iofile, read_dst, to_read, file_off); #endif io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); - /* Pack buffer index and staging slot into user_data for CQE dispatch */ - io_uring_sqe_set_data64(sqe, ((uint64_t)bidx << SHIFT_RANGE) | (uint64_t)slot); + io_uring_sqe_set_data64(sqe, ((uint64_t)bidx << FETCHV_BIDX_SHIFT) | (uint64_t)slot); st->submitted += to_read; n_in_flight++; @@ -383,32 +382,19 @@ static ssize_t hmll_io_uring_fetchv_impl( if (n_in_flight == 0) { if (n_active == 0) break; - - /* - * n_in_flight == 0 but n_active > 0: the inner loop couldn't claim - * any staging slot. For CUDA this means all 128 slots are occupied - * by pending cudaMemcpyAsync operations whose events haven't fired - * yet. Spinning here (nwait=0 path) would burn 100 % CPU for - * hundreds of milliseconds on large models. - * - * Fix: flush any queued fadvise SQEs, then block on the first - * outstanding CUDA event so we reclaim at least one slot before - * retrying. For CPU this case is transient (SQPOLL catching up), so - * the io_uring_submit() wake-up is sufficient. - */ io_uring_submit(&fetcher->ioring); - #if defined(__HMLL_CUDA_ENABLED__) - if (is_cuda && hmll_io_uring_slot_find_available(fetcher->iobusy) == -1) { + if (is_cuda && hmll_io_uring_slot_find_available(fetcher->iobusy) < 0) { struct hmll_io_uring_cuda_context *dctx = fetcher->device_ctx; - for (size_t i = 0; i < HMLL_URING_QUEUE_DEPTH; ++i) { - struct hmll_io_uring_cuda_context *cd = dctx + i; - if (hmll_io_uring_slot_is_busy(fetcher->iobusy, i) && - cd->state == HMLL_CUDA_STREAM_MEMCPY) { - cudaEventSynchronize(cd->done); - hmll_io_uring_cuda_stream_set_idle(&cd->state); - hmll_io_uring_slot_set_available(&fetcher->iobusy, i); - break; + for (size_t i = 0; i < HMLL_URING_QUEUE_DEPTH; i++) { + if (hmll_io_uring_slot_is_busy(fetcher->iobusy, i)) { + struct hmll_io_uring_cuda_context *cd = &dctx[i]; + if (cd->state == HMLL_CUDA_STREAM_MEMCPY) { + cudaEventSynchronize(cd->done); + hmll_io_uring_cuda_stream_set_idle(&cd->state); + hmll_io_uring_slot_set_available(&fetcher->iobusy, (unsigned)i); + break; + } } } } @@ -417,16 +403,13 @@ static ssize_t hmll_io_uring_fetchv_impl( } const size_t nwait = n_in_flight < fetcher->iocca.window ? n_in_flight : fetcher->iocca.window; - struct timespec ts_start, ts_end; clock_gettime(CLOCK_MONOTONIC, &ts_start); - - if (unlikely(io_uring_submit_and_wait(&fetcher->ioring, nwait) < 0)) { + if (io_uring_submit_and_wait(&fetcher->ioring, nwait) < 0) { ctx->error = HMLL_ERR(HMLL_ERR_IO_ERROR); goto cleanup; } clock_gettime(CLOCK_MONOTONIC, &ts_end); - hmll_io_uring_cca_update(&fetcher->iocca, HMLL_URING_BUFFER_SIZE * nwait, ts_start, ts_end); unsigned count; @@ -435,33 +418,28 @@ static ssize_t hmll_io_uring_fetchv_impl( const struct io_uring_cqe *cqe = cqes[i]; const uint64_t data = cqe->user_data; - /* Skip fadvise completions; use exact equality to avoid false - positives if future encodings ever set high bits. */ - if (unlikely(data == BIT_FADVISE)) continue; + if (data == FETCHV_FADVISE_TAG) continue; n_in_flight--; - - if (unlikely(cqe->res < 0)) { + if (cqe->res < 0) { ctx->error = HMLL_SYS_ERR(-cqe->res); io_uring_cq_advance(&fetcher->ioring, count); goto cleanup; } + nbytes += (size_t)cqe->res; - nbytes += cqe->res; - - const uint32_t slot = (uint32_t)(data & MASK_SLOT); - const uint32_t bidx = (uint32_t)(data >> SHIFT_RANGE); + const uint32_t slot = (uint32_t)(data & FETCHV_SLOT_MASK); + const uint32_t bidx = (uint32_t)(data >> FETCHV_BIDX_SHIFT); if (!is_cuda) { hmll_io_uring_slot_set_available(&fetcher->iobusy, slot); } #if defined(__HMLL_CUDA_ENABLED__) else { - struct hmll_io_uring_cuda_context *cctx = - &((struct hmll_io_uring_cuda_context *)fetcher->device_ctx)[slot]; + struct hmll_io_uring_cuda_context *cctx = &((struct hmll_io_uring_cuda_context *)fetcher->device_ctx)[slot]; + void *to = (char *)dsts[bidx].ptr + slot_offsets[slot]; void *from = fetcher->iovecs[slot].iov_base; - void *to = (char *)dsts[bidx].ptr + slot_offsets[slot]; - cudaMemcpyAsync(to, from, cqe->res, cudaMemcpyHostToDevice, cctx->stream); + cudaMemcpyAsync(to, from, (size_t)cqe->res, cudaMemcpyHostToDevice, cctx->stream); cudaEventRecord(cctx->done, cctx->stream); hmll_io_uring_cuda_stream_set_memcpy(&cctx->state); } @@ -474,11 +452,11 @@ static ssize_t hmll_io_uring_fetchv_impl( } hmll_io_uring_sync(dsts[0].device, fetcher); - if ((unsigned char *)states != stack_mem) free(states); + if (scratch_to_free) free(scratch_to_free); return (ssize_t)nbytes; cleanup: - if ((unsigned char *)states != stack_mem) free(states); + if (scratch_to_free) free(scratch_to_free); return -1; } From 510e4dfb53e76f0bbec76fd39e97cac28874132e Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 2 Mar 2026 14:55:02 +0100 Subject: [PATCH 05/10] add fetchv test suite to CTest --- CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index fc85808..9f1a18e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -232,6 +232,7 @@ if(${HMLL_BUILD_TESTS}) tests/tests_hmll_dtype.cpp tests/tests_hmll_safetensors.cpp tests/tests_hmll_integration_safetensors.cpp + tests/tests_hmll_fetchv.cpp ) endif() From f9ecb7bc12e969f9acbe387589b2009b9ffbd7f6 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 2 Mar 2026 15:04:58 +0100 Subject: [PATCH 06/10] add fetchv unittest --- tests/tests_hmll_fetchv.cpp | 612 ++++++++++++++++++++++++++++++++++++ 1 file changed, 612 insertions(+) create mode 100644 tests/tests_hmll_fetchv.cpp diff --git a/tests/tests_hmll_fetchv.cpp b/tests/tests_hmll_fetchv.cpp new file mode 100644 index 0000000..d964de6 --- /dev/null +++ b/tests/tests_hmll_fetchv.cpp @@ -0,0 +1,612 @@ +// +// Test suite for hmll_fetchv public API (mmap, io_uring backends). +// Uses fixtures from create_fetchv_testing_safetensors.py; set +// HMLL_CI_FETCHV_SAFETENSORS_FPATH and HMLL_CI_FETCHV_SHARDED_SAFETENSORS_FPATH. +// + +#include +#include "hmll/hmll.h" +#include +#include +#include +#include + +#if defined(__linux__) +#endif + +#define HMLL_CI_FETCHV_SAFETENSORS_FPATH "HMLL_CI_FETCHV_SAFETENSORS_FPATH" +#define HMLL_CI_FETCHV_SHARDED_SAFETENSORS_FPATH "HMLL_CI_FETCHV_SHARDED_SAFETENSORS_FPATH" + +namespace { + +using BackendPair = std::pair; +#if defined(__linux__) +constexpr BackendPair kBackends[] = { + {"IO_URING", HMLL_FETCHER_IO_URING}, + {"MMAP", HMLL_FETCHER_MMAP}, +}; +#else +constexpr BackendPair kBackends[] = {{"MMAP", HMLL_FETCHER_MMAP}}; +#endif + +// Validate float32 buffer filled with deterministic pattern value[i] = i (as float) +void validate_float32_arange(const hmll_iobuf_t& buf, size_t numel, size_t start = 0) { + REQUIRE(buf.ptr != nullptr); + REQUIRE(buf.size >= numel * sizeof(float)); + const auto* p = static_cast(buf.ptr); + for (size_t i = 0; i < numel; ++i) + REQUIRE(std::abs(p[i] - static_cast(start + i)) < 1e-5f); +} + +void validate_int32_deterministic(const hmll_iobuf_t& buf, size_t numel) { + REQUIRE(buf.ptr != nullptr); + REQUIRE(buf.size >= numel * sizeof(int32_t)); + const auto* p = static_cast(buf.ptr); + for (size_t i = 0; i < numel; ++i) { + const auto expected = (static_cast(i) % (1 << 15)) - (1 << 14); + REQUIRE(p[i] == static_cast(expected)); + } +} + +void validate_uint8_deterministic(const hmll_iobuf_t& buf, size_t numel) { + REQUIRE(buf.ptr != nullptr); + REQUIRE(buf.size >= numel); + const auto* p = static_cast(buf.ptr); + for (size_t i = 0; i < numel; ++i) + REQUIRE(p[i] == static_cast(i % 256)); +} + +// Return total bytes in dsts[0..n-1] +size_t total_dst_size(const hmll_iobuf_t* dsts, size_t n) { + size_t t = 0; + for (size_t i = 0; i < n; ++i) t += dsts[i].size; + return t; +} + +} // namespace + +TEST_CASE("fetchv - single-element fetchv (n=1)", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + + const auto* tensor_name = "float32.vec16"; + hmll_lookup_result_t lookup = hmll_lookup_tensor(&ctx, ®istry, tensor_name); + REQUIRE_FALSE(hmll_check(ctx.error)); + REQUIRE(lookup.specs != nullptr); + + hmll_range_t range = {lookup.specs->start, lookup.specs->end}; + hmll_iobuf_t buffer = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, range); + REQUIRE_FALSE(hmll_check(ctx.error)); + + size_t offsets[1] = {range.start}; + hmll_iobuf_t dsts[1] = {buffer}; + ssize_t ret = hmll_fetchv(&ctx, lookup.file, dsts, offsets, 1); + REQUIRE(ret >= 0); + REQUIRE(static_cast(ret) == buffer.size); + validate_float32_arange(buffer, 16); + + hmll_free_buffer(&buffer); + hmll_destroy(&ctx); + } + hmll_free_registry(®istry); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - multi-element same dtype", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + + const char* names[] = {"float32.vec16", "float32.vec1024", "float32.vec8192", "float32.scalar"}; + hmll_iobuf_t dsts[4]; + size_t offsets[4]; + size_t total = 0; + for (int i = 0; i < 4; ++i) { + hmll_lookup_result_t lookup = hmll_lookup_tensor(&ctx, ®istry, names[i]); + REQUIRE(lookup.specs != nullptr); + hmll_range_t range = {lookup.specs->start, lookup.specs->end}; + dsts[i] = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, range); + REQUIRE_FALSE(hmll_check(ctx.error)); + offsets[i] = range.start; + total += dsts[i].size; + } + + ssize_t ret = hmll_fetchv(&ctx, 0, dsts, offsets, 4); + REQUIRE(ret >= 0); + REQUIRE(static_cast(ret) == total); + validate_float32_arange(dsts[0], 16); + validate_float32_arange(dsts[1], 1024); + validate_float32_arange(dsts[2], 8192); + validate_float32_arange(dsts[3], 1); + + for (auto & dst : dsts) hmll_free_buffer(&dst); + hmll_destroy(&ctx); + } + hmll_free_registry(®istry); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - multi-element mixed dtypes", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + + hmll_lookup_result_t l_f32 = hmll_lookup_tensor(&ctx, ®istry, "float32.vec16"); + hmll_lookup_result_t l_i32 = hmll_lookup_tensor(&ctx, ®istry, "int32.vec16"); + hmll_lookup_result_t l_u8 = hmll_lookup_tensor(&ctx, ®istry, "uint8.vec16"); + REQUIRE((l_f32.specs && l_i32.specs && l_u8.specs)); + + hmll_iobuf_t dsts[3] = { + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, {l_f32.specs->start, l_f32.specs->end}), + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, {l_i32.specs->start, l_i32.specs->end}), + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, {l_u8.specs->start, l_u8.specs->end}), + }; + size_t offsets[3] = {l_f32.specs->start, l_i32.specs->start, l_u8.specs->start}; + REQUIRE_FALSE(hmll_check(ctx.error)); + + ssize_t ret = hmll_fetchv(&ctx, 0, dsts, offsets, 3); + REQUIRE(ret >= 0); + REQUIRE(static_cast(ret) == dsts[0].size + dsts[1].size + dsts[2].size); + validate_float32_arange(dsts[0], 16); + validate_int32_deterministic(dsts[1], 16); + validate_uint8_deterministic(dsts[2], 16); + + for (auto & dst : dsts) hmll_free_buffer(&dst); + hmll_destroy(&ctx); + } + hmll_free_registry(®istry); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - scattered reads within single tensor", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + + hmll_lookup_result_t lookup = hmll_lookup_tensor(&ctx, ®istry, "float32.vec8192"); + REQUIRE(lookup.specs != nullptr); + const size_t elem_size = sizeof(float); + const size_t base = lookup.specs->start; + + // Sub-ranges: [0..3], [100..103], [4000..4003] + struct { size_t start_off; size_t numel; } ranges[] = { + {0 * elem_size, 4}, + {100 * elem_size, 4}, + {4000 * elem_size, 4}, + }; + hmll_iobuf_t dsts[3]; + size_t offsets[3]; + for (int i = 0; i < 3; ++i) { + size_t len = ranges[i].numel * elem_size; + hmll_range_t r = {base + ranges[i].start_off, base + ranges[i].start_off + len}; + dsts[i] = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, r); + offsets[i] = r.start; + } + + ssize_t ret = hmll_fetchv(&ctx, lookup.file, dsts, offsets, 3); + REQUIRE(ret >= 0); + validate_float32_arange(dsts[0], 4, 0); + validate_float32_arange(dsts[1], 4, 100); + validate_float32_arange(dsts[2], 4, 4000); + + for (auto & dst : dsts) hmll_free_buffer(&dst); + hmll_destroy(&ctx); + } + hmll_free_registry(®istry); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - full tensor via fetchv matches hmll_fetch", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + + hmll_lookup_result_t lookup = hmll_lookup_tensor(&ctx, ®istry, "float32.vec1024"); + REQUIRE(lookup.specs != nullptr); + hmll_range_t range = {lookup.specs->start, lookup.specs->end}; + hmll_iobuf_t buf_fetch = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, range); + hmll_iobuf_t buf_fetchv = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, range); + REQUIRE_FALSE(hmll_check(ctx.error)); + + ssize_t r1 = hmll_fetch(&ctx, lookup.file, &buf_fetch, range.start); + size_t offsets[1] = {range.start}; + hmll_iobuf_t dsts[1] = {buf_fetchv}; + ssize_t r2 = hmll_fetchv(&ctx, lookup.file, dsts, offsets, 1); + REQUIRE(r1 > 0); + REQUIRE(r2 > 0); + REQUIRE(static_cast(r1) == buf_fetch.size); + REQUIRE(static_cast(r2) == buf_fetchv.size); + REQUIRE(std::memcmp(buf_fetch.ptr, buf_fetchv.ptr, buf_fetch.size) == 0); + + hmll_free_buffer(&buf_fetch); + hmll_free_buffer(&buf_fetchv); + hmll_destroy(&ctx); + } + hmll_free_registry(®istry); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - n=0 returns 0", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, kBackends[0].second))); + + ssize_t ret = hmll_fetchv(&ctx, 0, nullptr, nullptr, 0); + REQUIRE(ret == 0); + REQUIRE_FALSE(hmll_check(ctx.error)); + + hmll_free_registry(®istry); + hmll_destroy(&ctx); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - scalar tensors", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + + hmll_lookup_result_t l0 = hmll_lookup_tensor(&ctx, ®istry, "float32.scalar"); + hmll_lookup_result_t l1 = hmll_lookup_tensor(&ctx, ®istry, "int32.scalar"); + hmll_lookup_result_t l2 = hmll_lookup_tensor(&ctx, ®istry, "uint8.scalar"); + REQUIRE((l0.specs && l1.specs && l2.specs)); + + hmll_iobuf_t dsts[3] = { + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, {l0.specs->start, l0.specs->end}), + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, {l1.specs->start, l1.specs->end}), + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, {l2.specs->start, l2.specs->end}), + }; + size_t offsets[3] = {l0.specs->start, l1.specs->start, l2.specs->start}; + + ssize_t ret = hmll_fetchv(&ctx, 0, dsts, offsets, 3); + REQUIRE(ret >= 0); + REQUIRE(static_cast(ret) == dsts[0].size + dsts[1].size + dsts[2].size); + validate_float32_arange(dsts[0], 1); + validate_int32_deterministic(dsts[1], 1); + validate_uint8_deterministic(dsts[2], 1); + + for (auto & dst : dsts) hmll_free_buffer(&dst); + hmll_destroy(&ctx); + } + hmll_free_registry(®istry); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - large tensor exceeds io_uring buffer", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + + hmll_lookup_result_t lookup = hmll_lookup_tensor(&ctx, ®istry, "float32.large"); + REQUIRE(lookup.specs != nullptr); + hmll_range_t range = {lookup.specs->start, lookup.specs->end}; + hmll_iobuf_t buffer = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, range); + REQUIRE_FALSE(hmll_check(ctx.error)); + REQUIRE(buffer.size > 512 * 1024u); // > HMLL_URING_BUFFER_SIZE + + size_t offsets[1] = {range.start}; + hmll_iobuf_t dsts[1] = {buffer}; + ssize_t ret = hmll_fetchv(&ctx, lookup.file, dsts, offsets, 1); + REQUIRE(ret >= 0); + REQUIRE(static_cast(ret) == buffer.size); + size_t numel = buffer.size / sizeof(float); + validate_float32_arange(buffer, numel); + + hmll_free_buffer(&buffer); + hmll_destroy(&ctx); + } + hmll_free_registry(®istry); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - many concurrent ranges", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + + const size_t N = 40; + std::vector lookups(N); + std::vector dsts(N); + std::vector offsets(N); + for (size_t i = 0; i < N; ++i) { + lookups[i] = hmll_lookup_tensor(&ctx, ®istry, "float32.vec16"); + REQUIRE(lookups[i].specs != nullptr); + hmll_range_t range = {lookups[i].specs->start, lookups[i].specs->end}; + dsts[i] = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, range); + offsets[i] = range.start; + } + + ssize_t ret = hmll_fetchv(&ctx, 0, dsts.data(), offsets.data(), N); + REQUIRE(ret >= 0); + REQUIRE(static_cast(ret) == total_dst_size(dsts.data(), N)); + for (size_t i = 0; i < N; ++i) + validate_float32_arange(dsts[i], 16); + + for (size_t i = 0; i < N; ++i) hmll_free_buffer(&dsts[i]); + hmll_destroy(&ctx); + } + hmll_free_registry(®istry); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - return value equals sum of dst sizes", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, kBackends[0].second))); + + const char* names[] = {"float32.vec16", "int32.vec16", "float32.scalar"}; + hmll_iobuf_t dsts[3]; + size_t offsets[3]; + size_t expected_total = 0; + for (int i = 0; i < 3; ++i) { + hmll_lookup_result_t l = hmll_lookup_tensor(&ctx, ®istry, names[i]); + REQUIRE(l.specs != nullptr); + hmll_range_t r = {l.specs->start, l.specs->end}; + dsts[i] = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, r); + offsets[i] = r.start; + expected_total += dsts[i].size; + } + ssize_t ret = hmll_fetchv(&ctx, 0, dsts, offsets, 3); + REQUIRE(ret >= 0); + REQUIRE(static_cast(ret) == expected_total); + + for (auto & dst : dsts) hmll_free_buffer(&dst); + hmll_free_registry(®istry); + hmll_destroy(&ctx); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - overlapping logical range same data", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, kBackends[0].second))); + + hmll_lookup_result_t lookup = hmll_lookup_tensor(&ctx, ®istry, "float32.vec16"); + REQUIRE(lookup.specs != nullptr); + hmll_range_t range = {lookup.specs->start, lookup.specs->end}; + hmll_iobuf_t buf1 = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, range); + hmll_iobuf_t buf2 = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, range); + size_t offsets[2] = {range.start, range.start}; + hmll_iobuf_t dsts[2] = {buf1, buf2}; + ssize_t ret = hmll_fetchv(&ctx, lookup.file, dsts, offsets, 2); + REQUIRE(ret >= 0); + REQUIRE(std::memcmp(buf1.ptr, buf2.ptr, buf1.size) == 0); + hmll_free_buffer(&buf1); + hmll_free_buffer(&buf2); + hmll_free_registry(®istry); + hmll_destroy(&ctx); + hmll_source_close(&src); +} + +// --- Sharded tests --- +TEST_CASE("fetchv - sharded index parsing", "[fetchv][safetensors][sharded]") { + const char* index_path = std::getenv(HMLL_CI_FETCHV_SHARDED_SAFETENSORS_FPATH); + if (!index_path) SKIP("HMLL_CI_FETCHV_SHARDED_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t index_src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(index_path, &index_src))); + hmll_registry_t registry = {}; + size_t num_files = hmll_safetensors_index(&ctx, ®istry, index_src); + REQUIRE(num_files > 0); + REQUIRE(registry.num_tensors > 0); + REQUIRE(registry.indexes != nullptr); + REQUIRE(registry.names != nullptr); + REQUIRE(registry.tensors != nullptr); + hmll_source_close(&index_src); + hmll_free_registry(®istry); +} + +TEST_CASE("fetchv - sharded fetchv across files", "[fetchv][safetensors][sharded]") { + const char* index_path = std::getenv(HMLL_CI_FETCHV_SHARDED_SAFETENSORS_FPATH); + if (!index_path) SKIP("HMLL_CI_FETCHV_SHARDED_SAFETENSORS_FPATH not set"); + + std::string dir(index_path); + size_t slash = dir.find_last_of("/\\"); + if (slash != std::string::npos) dir.resize(slash + 1); + else dir = "./"; + + hmll_t ctx = {}; + hmll_source_t index_src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(index_path, &index_src))); + hmll_registry_t registry = {}; + size_t num_files = hmll_safetensors_index(&ctx, ®istry, index_src); + REQUIRE(num_files > 0); + hmll_source_close(&index_src); + + std::vector sources(num_files); + for (size_t i = 0; i < num_files; ++i) { + char path[512]; + int n = std::snprintf(path, sizeof(path), "%smodel-%05zu-of-%05zu.safetensors", + dir.c_str(), i + 1, num_files); + REQUIRE((n > 0 && static_cast(n) < sizeof(path))); + REQUIRE_FALSE(hmll_check(hmll_source_open(path, &sources[i]))); + } + + size_t offset = 0; + for (size_t i = 0; i < num_files; ++i) { + size_t n = hmll_safetensors_populate_registry(&ctx, ®istry, sources[i], (unsigned short)i, offset); + REQUIRE(n > 0); + offset += n; + } + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, sources.data(), num_files, HMLL_DEVICE_CPU, kBackends[0].second))); + + for (const auto& [name, backend] : kBackends) { + hmll_destroy(&ctx); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, sources.data(), num_files, HMLL_DEVICE_CPU, backend))); + + hmll_lookup_result_t l0 = hmll_lookup_tensor(&ctx, ®istry, "float32.shard0.vec16"); + hmll_lookup_result_t l1 = hmll_lookup_tensor(&ctx, ®istry, "int32.shard1.vec16"); + hmll_lookup_result_t l2 = hmll_lookup_tensor(&ctx, ®istry, "bfloat16.shard2.vec64"); + REQUIRE(l0.specs != nullptr); + REQUIRE(l1.specs != nullptr); + REQUIRE(l2.specs != nullptr); + REQUIRE(l0.file == 0); + REQUIRE(l1.file == 1); + REQUIRE(l2.file == 2); + + hmll_iobuf_t dsts[3] = { + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, {l0.specs->start, l0.specs->end}), + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, {l1.specs->start, l1.specs->end}), + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, {l2.specs->start, l2.specs->end}), + }; + size_t offsets[3] = {l0.specs->start, l1.specs->start, l2.specs->start}; + int iofiles[3] = {l0.file, l1.file, l2.file}; + + ssize_t total = 0; + for (int i = 0; i < 3; ++i) { + hmll_iobuf_t d[1] = {dsts[i]}; + size_t o[1] = {offsets[i]}; + ssize_t r = hmll_fetchv(&ctx, iofiles[i], d, o, 1); + REQUIRE(r >= 0); + total += r; + } + validate_float32_arange(dsts[0], 16); + validate_int32_deterministic(dsts[1], 16); + REQUIRE(dsts[2].size == 64 * 2u); // bf16 + for (auto & dst : dsts) hmll_free_buffer(&dst); + } + + for (size_t i = 0; i < num_files; ++i) hmll_source_close(&sources[i]); + hmll_free_registry(®istry); + hmll_destroy(&ctx); +} + +TEST_CASE("fetchv - fetchv matches fetch for several tensors", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + + const char* names[] = {"float32.vec16", "int32.vec16", "uint8.vec16", "float32.scalar"}; + const size_t n_tensors = 4; + std::vector buf_fetch(n_tensors); + std::vector buf_fetchv(n_tensors); + std::vector offsets(n_tensors); + size_t total = 0; + for (size_t i = 0; i < n_tensors; ++i) { + hmll_lookup_result_t l = hmll_lookup_tensor(&ctx, ®istry, names[i]); + REQUIRE(l.specs != nullptr); + hmll_range_t r = {l.specs->start, l.specs->end}; + buf_fetch[i] = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, r); + buf_fetchv[i] = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, r); + offsets[i] = r.start; + total += buf_fetch[i].size; + } + for (size_t i = 0; i < n_tensors; ++i) { + hmll_lookup_result_t l = hmll_lookup_tensor(&ctx, ®istry, names[i]); + ssize_t r = hmll_fetch(&ctx, l.file, &buf_fetch[i], offsets[i]); + REQUIRE(r > 0); + } + ssize_t rv = hmll_fetchv(&ctx, 0, buf_fetchv.data(), offsets.data(), n_tensors); + REQUIRE(rv >= 0); + REQUIRE(static_cast(rv) == total); + for (size_t i = 0; i < n_tensors; ++i) + REQUIRE(std::memcmp(buf_fetch[i].ptr, buf_fetchv[i].ptr, buf_fetch[i].size) == 0); + for (size_t i = 0; i < n_tensors; ++i) { + hmll_free_buffer(&buf_fetch[i]); + hmll_free_buffer(&buf_fetchv[i]); + } + hmll_destroy(&ctx); + } + hmll_free_registry(®istry); + hmll_source_close(&src); +} From ad5a19176110c2ac27e08286dd32089c2fd86b7b Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 2 Mar 2026 22:08:37 +0100 Subject: [PATCH 07/10] uniformize iouring fadvise flag between fetch/fetchv --- lib/linux/backend/iouring.c | 11 ++- tests/tests_hmll_fetchv.cpp | 137 ++++++++++++++++++++++++++++++++++++ 2 files changed, 142 insertions(+), 6 deletions(-) diff --git a/lib/linux/backend/iouring.c b/lib/linux/backend/iouring.c index fbf72b8..5100be4 100644 --- a/lib/linux/backend/iouring.c +++ b/lib/linux/backend/iouring.c @@ -4,7 +4,7 @@ #include "hmll/linux/backend/iouring.h" #include "sys/mman.h" -#define HMLL_IO_URING_ADVISORY_FLAG UINT64_MAX +#define HMLL_IO_URING_FADVISE_TAG (1ULL << 63) #if defined(__HMLL_CUDA_ENABLED__) #include "hmll/cuda.h" @@ -194,7 +194,7 @@ static ssize_t hmll_io_uring_fetch_impl( if ((sqe = io_uring_get_sqe(&fetcher->ioring))) { io_uring_prep_fadvise(sqe, iofile, offset, dst->size, POSIX_FADV_WILLNEED); io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); - io_uring_sqe_set_data64(sqe, HMLL_IO_URING_ADVISORY_FLAG); + io_uring_sqe_set_data64(sqe, HMLL_IO_URING_FADVISE_TAG); } while (b_read < dst->size) { @@ -237,7 +237,7 @@ static ssize_t hmll_io_uring_fetch_impl( for (unsigned i = 0; i < count; i++) { const struct io_uring_cqe *cqe = cqes[i]; - if (unlikely(cqe->user_data == HMLL_IO_URING_ADVISORY_FLAG)) + if (unlikely(cqe->user_data == HMLL_IO_URING_FADVISE_TAG)) continue; --n_dma; @@ -273,7 +273,6 @@ static ssize_t hmll_io_uring_fetchv_impl( const int is_cuda = (dsts[0].device == HMLL_DEVICE_CUDA); /* user_data encoding for CQEs: high bit = fadvise (skip), else (bidx << 8) | slot */ - static const uint64_t FETCHV_FADVISE_TAG = 1ULL << 63; static const unsigned FETCHV_BIDX_SHIFT = 8; static const uint64_t FETCHV_SLOT_MASK = HMLL_URING_QUEUE_DEPTH - 1; @@ -340,7 +339,7 @@ static ssize_t hmll_io_uring_fetchv_impl( if (!sqe) break; io_uring_prep_fadvise(sqe, iofile, offsets[bidx], st->size, POSIX_FADV_WILLNEED); io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); - io_uring_sqe_set_data64(sqe, FETCHV_FADVISE_TAG); + io_uring_sqe_set_data64(sqe, HMLL_IO_URING_FADVISE_TAG); st->fadvise_sent = 1; continue; } @@ -418,7 +417,7 @@ static ssize_t hmll_io_uring_fetchv_impl( const struct io_uring_cqe *cqe = cqes[i]; const uint64_t data = cqe->user_data; - if (data == FETCHV_FADVISE_TAG) continue; + if (data == HMLL_IO_URING_FADVISE_TAG) continue; n_in_flight--; if (cqe->res < 0) { diff --git a/tests/tests_hmll_fetchv.cpp b/tests/tests_hmll_fetchv.cpp index d964de6..2af8607 100644 --- a/tests/tests_hmll_fetchv.cpp +++ b/tests/tests_hmll_fetchv.cpp @@ -610,3 +610,140 @@ TEST_CASE("fetchv - fetchv matches fetch for several tensors", "[fetchv][safeten hmll_free_registry(®istry); hmll_source_close(&src); } + +TEST_CASE("fetchv - pre-existing error returns -1", "[fetchv][error]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, kBackends[0].second))); + + // Poison the context with an error + ctx.error = HMLL_ERR(HMLL_ERR_IO_ERROR); + + hmll_iobuf_t dsts[1] = {}; + size_t offsets[1] = {0}; + ssize_t ret = hmll_fetchv(&ctx, 0, dsts, offsets, 1); + REQUIRE(ret == -1); + + ctx.error = HMLL_OK; // reset so cleanup works + hmll_free_registry(®istry); + hmll_destroy(&ctx); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - multiple large tensors interleaved", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + + // Two large tensors, each > HMLL_URING_BUFFER_SIZE + const char* names[] = {"float32.large", "int32.large"}; + hmll_iobuf_t dsts[2]; + size_t offsets[2]; + size_t total = 0; + for (int i = 0; i < 2; ++i) { + hmll_lookup_result_t l = hmll_lookup_tensor(&ctx, ®istry, names[i]); + REQUIRE(l.specs != nullptr); + hmll_range_t r = {l.specs->start, l.specs->end}; + dsts[i] = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, r); + offsets[i] = r.start; + total += dsts[i].size; + REQUIRE(dsts[i].size > 512 * 1024u); + } + + ssize_t ret = hmll_fetchv(&ctx, 0, dsts, offsets, 2); + REQUIRE(ret >= 0); + REQUIRE(static_cast(ret) == total); + validate_float32_arange(dsts[0], dsts[0].size / sizeof(float)); + validate_int32_deterministic(dsts[1], dsts[1].size / sizeof(int32_t)); + + for (auto& dst : dsts) hmll_free_buffer(&dst); + hmll_destroy(&ctx); + } + hmll_free_registry(®istry); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - heap scratch path (large N)", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, kBackends[0].second))); + + const size_t N = 300; // exceeds stack_scratch[8192] + hmll_lookup_result_t l = hmll_lookup_tensor(&ctx, ®istry, "float32.vec16"); + REQUIRE(l.specs != nullptr); + + std::vector dsts(N); + std::vector offsets(N); + for (size_t i = 0; i < N; ++i) { + hmll_range_t r = {l.specs->start, l.specs->end}; + dsts[i] = hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, r); + offsets[i] = r.start; + } + + ssize_t ret = hmll_fetchv(&ctx, 0, dsts.data(), offsets.data(), N); + REQUIRE(ret >= 0); + REQUIRE(static_cast(ret) == total_dst_size(dsts.data(), N)); + for (size_t i = 0; i < N; ++i) + validate_float32_arange(dsts[i], 16); + + for (size_t i = 0; i < N; ++i) hmll_free_buffer(&dsts[i]); + hmll_free_registry(®istry); + hmll_destroy(&ctx); + hmll_source_close(&src); +} + +TEST_CASE("fetchv - zero-size buffer interspersed", "[fetchv][safetensors]") { + const char* fpath = std::getenv(HMLL_CI_FETCHV_SAFETENSORS_FPATH); + if (!fpath) SKIP("HMLL_CI_FETCHV_SAFETENSORS_FPATH not set"); + + hmll_t ctx = {}; + hmll_source_t src = {}; + REQUIRE_FALSE(hmll_check(hmll_source_open(fpath, &src))); + hmll_registry_t registry = {}; + REQUIRE(hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) > 0); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, kBackends[0].second))); + + hmll_lookup_result_t l = hmll_lookup_tensor(&ctx, ®istry, "float32.vec16"); + REQUIRE(l.specs != nullptr); + hmll_range_t r = {l.specs->start, l.specs->end}; + + hmll_iobuf_t dsts[3] = { + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, r), + {.size = 0, .ptr = nullptr, .device = HMLL_DEVICE_CPU}, // zero-size + hmll_get_buffer_for_range(&ctx, ctx.fetcher->device, r), + }; + size_t offsets[3] = {r.start, 0, r.start}; + + ssize_t ret = hmll_fetchv(&ctx, 0, dsts, offsets, 3); + REQUIRE(ret >= 0); + REQUIRE(static_cast(ret) == dsts[0].size + dsts[2].size); + validate_float32_arange(dsts[0], 16); + validate_float32_arange(dsts[2], 16); + + hmll_free_buffer(&dsts[0]); + hmll_free_buffer(&dsts[2]); + hmll_free_registry(®istry); + hmll_destroy(&ctx); + hmll_source_close(&src); +} \ No newline at end of file From 8f2aa659bc35c7f655eb44e0bed62c3c8a2a6441 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Wed, 4 Mar 2026 15:51:23 +0100 Subject: [PATCH 08/10] improve on the capacity guard for safetensors binding --- include/hmll/hmll.h | 1 + lib/python/loader.cpp | 6 +++--- lib/python/loader.hpp | 4 ++-- lib/python/safetensors.cpp | 12 ++++++++++-- lib/tensors.c | 14 ++++++++++++-- 5 files changed, 28 insertions(+), 9 deletions(-) diff --git a/include/hmll/hmll.h b/include/hmll/hmll.h index 124d82f..ad253e8 100644 --- a/include/hmll/hmll.h +++ b/include/hmll/hmll.h @@ -132,6 +132,7 @@ HMLL_EXTERN struct hmll_error hmll_get_mmap_view(struct hmll *ctx, int iofile, s /** Tensors manipulation stubs - enabled if a higher-level tensor format is enabled **/ #ifdef __HMLL_TENSORS_ENABLED__ HMLL_EXTERN uint8_t hmll_nbits(enum hmll_dtype dtype) NO_EXCEPT; +HMLL_EXTERN size_t hmll_nbytes(const struct hmll_tensor_specs *specs) NO_EXCEPT; HMLL_EXTERN size_t hmll_numel(const struct hmll_tensor_specs *specs) NO_EXCEPT; HMLL_EXTERN void hmll_free_registry(struct hmll_registry *reg) NO_EXCEPT; HMLL_EXTERN unsigned char hmll_contains(const struct hmll *ctx, const struct hmll_registry *reg, const char *name) NO_EXCEPT; diff --git a/lib/python/loader.cpp b/lib/python/loader.cpp index 864f59e..530405d 100644 --- a/lib/python/loader.cpp +++ b/lib/python/loader.cpp @@ -78,15 +78,15 @@ nb::ndarray WeightLoader::afetch(const int iofile, const size_t st return hmll_to_ndarray(buffer, dtype, shape, rank, std::move(deleter)); } -size_t WeightLoader::fetch(const int iofile, const size_t offset, const uintptr_t dst, const size_t size) const +size_t WeightLoader::fetch(const int iofile, const size_t offset, void *const dst, const size_t capacity) const { nb::gil_scoped_release release; const auto ctx = ctx_.get(); const auto dev = device(); - const hmll_iobuf_t buf = {size, reinterpret_cast(dst), dev}; - const auto [start, end] = hmll_range_t{offset, offset + size}; + const hmll_iobuf_t buf = {capacity, dst, dev}; + const auto [start, end] = hmll_range_t{offset, offset + capacity}; if (const auto res = hmll_fetch(ctx, iofile, &buf, start); res <= 0) { const std::string err = hmll_strerr(ctx_->error); throw std::runtime_error(fmt::format(PYHMLL_ERR_FETCH, err)); diff --git a/lib/python/loader.hpp b/lib/python/loader.hpp index 23bbaa0..55c9574 100644 --- a/lib/python/loader.hpp +++ b/lib/python/loader.hpp @@ -42,9 +42,9 @@ class WeightLoader afetch(int iofile, size_t start, size_t end, hmll_dtype_t dtype, const size_t* shape, uint8_t rank) const; [[nodiscard]] size_t - fetch(int iofile, size_t offset, uintptr_t dst, size_t size) const; + fetch(int iofile, size_t offset, void *dst, size_t capacity) const; - /** Batched fetch: write multiple byte ranges from file into a single pre-allocated buffer. */ + /** Batched fetch: write multiple byte ranges from the file into a single pre-allocated buffer. */ [[nodiscard]] size_t fetchv(int iofile, const std::vector>& ranges, uintptr_t dst) const; }; diff --git a/lib/python/safetensors.cpp b/lib/python/safetensors.cpp index aa8ceab..883b497 100644 --- a/lib/python/safetensors.cpp +++ b/lib/python/safetensors.cpp @@ -142,7 +142,7 @@ class SafetensorsAccessor return loader_->afetch(iofile, start, end, dtype, shape, rank); } - [[nodiscard]] size_t fetch(const std::string& name, const uintptr_t dst, const size_t size) const + [[nodiscard]] size_t fetch(const std::string& name, void *const dst, const size_t capacity) const { const auto registry = registry_.get(); const auto index = hmll_find_by_name(loader_->context(), registry, name.c_str()); @@ -152,9 +152,17 @@ class SafetensorsAccessor const auto specs = registry->tensors[index]; const auto iofile = registry->indexes[index]; + const auto nbytes = hmll_nbytes(&specs); + + if (capacity == 0) + throw std::runtime_error("Invalid 0-size for fetch operation"); + + if (capacity < nbytes) + throw std::runtime_error(fmt::format( + FMT_COMPILE("Provided destination buffer cannot be smaller than tensor size (provided={}, required={})"), capacity, nbytes)); // Delegate to WeightLoader for actual fetching - return loader_->fetch(iofile, specs.start, dst, size); + return loader_->fetch(iofile, specs.start, dst, capacity); } /** Fetch only the given element ranges into dst. Dtype is taken from the registry. */ diff --git a/lib/tensors.c b/lib/tensors.c index 0c430ce..5c96cc2 100644 --- a/lib/tensors.c +++ b/lib/tensors.c @@ -99,8 +99,18 @@ uint8_t hmll_nbits(const enum hmll_dtype dtype) } } -size_t hmll_numel(const hmll_tensor_specs_t *specs) +size_t hmll_nbytes(const struct hmll_tensor_specs *specs) { + if (!specs) return 0; + + const size_t numel = hmll_numel(specs); + const size_t nbits = hmll_nbits(specs->dtype); + return numel * (nbits / 8); +} + +size_t hmll_numel(const struct hmll_tensor_specs *specs) +{ + if (!specs) return 0; if (specs->rank > HMLL_MAX_TENSOR_RANK) __builtin_unreachable(); size_t numel = 1; @@ -108,4 +118,4 @@ size_t hmll_numel(const hmll_tensor_specs_t *specs) numel *= specs->shape[i]; return numel; -} +} \ No newline at end of file From 31f2b81a33b58ce7d57d9f157520de0b5edfe9d6 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Wed, 4 Mar 2026 16:03:21 +0100 Subject: [PATCH 09/10] ok, false good idea --- lib/python/loader.cpp | 4 ++-- lib/python/loader.hpp | 2 +- lib/python/safetensors.cpp | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/lib/python/loader.cpp b/lib/python/loader.cpp index 530405d..98bb978 100644 --- a/lib/python/loader.cpp +++ b/lib/python/loader.cpp @@ -78,14 +78,14 @@ nb::ndarray WeightLoader::afetch(const int iofile, const size_t st return hmll_to_ndarray(buffer, dtype, shape, rank, std::move(deleter)); } -size_t WeightLoader::fetch(const int iofile, const size_t offset, void *const dst, const size_t capacity) const +size_t WeightLoader::fetch(const int iofile, const size_t offset, const uintptr_t dst, const size_t capacity) const { nb::gil_scoped_release release; const auto ctx = ctx_.get(); const auto dev = device(); - const hmll_iobuf_t buf = {capacity, dst, dev}; + const hmll_iobuf_t buf = {capacity, reinterpret_cast(dst), dev}; const auto [start, end] = hmll_range_t{offset, offset + capacity}; if (const auto res = hmll_fetch(ctx, iofile, &buf, start); res <= 0) { const std::string err = hmll_strerr(ctx_->error); diff --git a/lib/python/loader.hpp b/lib/python/loader.hpp index 55c9574..f50e80a 100644 --- a/lib/python/loader.hpp +++ b/lib/python/loader.hpp @@ -42,7 +42,7 @@ class WeightLoader afetch(int iofile, size_t start, size_t end, hmll_dtype_t dtype, const size_t* shape, uint8_t rank) const; [[nodiscard]] size_t - fetch(int iofile, size_t offset, void *dst, size_t capacity) const; + fetch(int iofile, size_t offset, uintptr_t dst, size_t capacity) const; /** Batched fetch: write multiple byte ranges from the file into a single pre-allocated buffer. */ [[nodiscard]] size_t diff --git a/lib/python/safetensors.cpp b/lib/python/safetensors.cpp index 883b497..a885f79 100644 --- a/lib/python/safetensors.cpp +++ b/lib/python/safetensors.cpp @@ -142,7 +142,7 @@ class SafetensorsAccessor return loader_->afetch(iofile, start, end, dtype, shape, rank); } - [[nodiscard]] size_t fetch(const std::string& name, void *const dst, const size_t capacity) const + [[nodiscard]] size_t fetch(const std::string& name, const uintptr_t dst, const size_t capacity) const { const auto registry = registry_.get(); const auto index = hmll_find_by_name(loader_->context(), registry, name.c_str()); From b3cb8b693b2f335fd745bf1af63d8f414847469b Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Wed, 4 Mar 2026 16:15:23 +0100 Subject: [PATCH 10/10] pass underlying buffer size along pointer on Python binding --- lib/python/loader.cpp | 8 +++++--- lib/python/loader.hpp | 2 +- lib/python/safetensors.cpp | 12 +++++------- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/lib/python/loader.cpp b/lib/python/loader.cpp index 98bb978..baab8d5 100644 --- a/lib/python/loader.cpp +++ b/lib/python/loader.cpp @@ -78,15 +78,17 @@ nb::ndarray WeightLoader::afetch(const int iofile, const size_t st return hmll_to_ndarray(buffer, dtype, shape, rank, std::move(deleter)); } -size_t WeightLoader::fetch(const int iofile, const size_t offset, const uintptr_t dst, const size_t capacity) const +size_t WeightLoader::fetch(const int iofile, const size_t offset, const uintptr_t dst, const size_t size) const { nb::gil_scoped_release release; + if (size == 0) return 0; + const auto ctx = ctx_.get(); const auto dev = device(); - const hmll_iobuf_t buf = {capacity, reinterpret_cast(dst), dev}; - const auto [start, end] = hmll_range_t{offset, offset + capacity}; + const hmll_iobuf_t buf = {size, reinterpret_cast(dst), dev}; + const auto [start, end] = hmll_range_t{offset, offset + size}; if (const auto res = hmll_fetch(ctx, iofile, &buf, start); res <= 0) { const std::string err = hmll_strerr(ctx_->error); throw std::runtime_error(fmt::format(PYHMLL_ERR_FETCH, err)); diff --git a/lib/python/loader.hpp b/lib/python/loader.hpp index f50e80a..cafe076 100644 --- a/lib/python/loader.hpp +++ b/lib/python/loader.hpp @@ -42,7 +42,7 @@ class WeightLoader afetch(int iofile, size_t start, size_t end, hmll_dtype_t dtype, const size_t* shape, uint8_t rank) const; [[nodiscard]] size_t - fetch(int iofile, size_t offset, uintptr_t dst, size_t capacity) const; + fetch(int iofile, size_t offset, uintptr_t dst, size_t size) const; /** Batched fetch: write multiple byte ranges from the file into a single pre-allocated buffer. */ [[nodiscard]] size_t diff --git a/lib/python/safetensors.cpp b/lib/python/safetensors.cpp index a885f79..6d2ffba 100644 --- a/lib/python/safetensors.cpp +++ b/lib/python/safetensors.cpp @@ -142,7 +142,7 @@ class SafetensorsAccessor return loader_->afetch(iofile, start, end, dtype, shape, rank); } - [[nodiscard]] size_t fetch(const std::string& name, const uintptr_t dst, const size_t capacity) const + [[nodiscard]] size_t fetch(const std::string& name, const uintptr_t dst, const size_t size) const { const auto registry = registry_.get(); const auto index = hmll_find_by_name(loader_->context(), registry, name.c_str()); @@ -154,15 +154,13 @@ class SafetensorsAccessor const auto iofile = registry->indexes[index]; const auto nbytes = hmll_nbytes(&specs); - if (capacity == 0) - throw std::runtime_error("Invalid 0-size for fetch operation"); - - if (capacity < nbytes) + if (size < nbytes) throw std::runtime_error(fmt::format( - FMT_COMPILE("Provided destination buffer cannot be smaller than tensor size (provided={}, required={})"), capacity, nbytes)); + FMT_COMPILE("Provided destination buffer cannot be smaller than tensor size (provided={}, required={})"), size, nbytes)); + // Delegate to WeightLoader for actual fetching - return loader_->fetch(iofile, specs.start, dst, capacity); + return loader_->fetch(iofile, specs.start, dst, size); } /** Fetch only the given element ranges into dst. Dtype is taken from the registry. */