diff --git a/examples/fetch.c b/examples/fetch.c index 4295dc2..6c1c488 100644 --- a/examples/fetch.c +++ b/examples/fetch.c @@ -29,7 +29,7 @@ int main(const int argc, const char** argv) status = 3; goto exit; } - printf("\t- %s -> %zu bytes (fd: %u)\n", argv[i + 1], src[i].size, src[i].fd); + printf("\t- %s -> %zu bytes (b_fd: %d, d_fd: %d)\n", argv[i + 1], src[i].size, src[i].fd, src[i].d_fd); } err = hmll_loader_init(&ctx, src, argc - 1, hmll_device_cuda(0), HMLL_FETCHER_AUTO); diff --git a/examples/fetchv.c b/examples/fetchv.c index 0531721..55768cd 100644 --- a/examples/fetchv.c +++ b/examples/fetchv.c @@ -1,5 +1,6 @@ #include #include +#include #include #ifdef _WIN32 @@ -32,92 +33,249 @@ static double time_diff_ns(const timespec_t *start, const timespec_t *end) { #include #endif -// #define TENSOR_NAME "language_model.model.layers.15.mlp.gate_proj.weight" -#define TENSOR_NAME "float32.dim5" +static int path_ends_with(const char *path, const char *suffix) { + const size_t plen = strlen(path); + const size_t slen = strlen(suffix); + return plen >= slen && strcmp(path + plen - slen, suffix) == 0; +} + +static void get_dir_prefix(const char *path, char *dir, size_t dir_size) { + strncpy(dir, path, dir_size - 1); + dir[dir_size - 1] = '\0'; + char *last_slash = strrchr(dir, '/'); + if (last_slash) { + *(last_slash + 1) = '\0'; + } else { + dir[0] = '\0'; + } +} int main(const int argc, const char** argv) { if (argc < 2) { - printf("No file specified.\nInvoke through hmll_safetensors_ex "); + fprintf(stderr, "Usage: hmll_examples_fetchv \n"); return 1; } - hmll_t ctx = {0}; - hmll_source_t src = {0}; - if (hmll_check(hmll_source_open(argv[1], &src))) - return 1; + const char *path = argv[1]; + const int is_sharded = path_ends_with(path, ".index.json"); - // Read safetensors table with all tensors mapping + hmll_t ctx = {0}; hmll_registry_t registry = {0}; - if (hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) == 0) - return 2; + hmll_source_t *sources = NULL; + size_t num_files = 1; + + if (is_sharded) { + hmll_source_t index_src = {0}; + if (hmll_check(hmll_source_open(path, &index_src))) { + fprintf(stderr, "Failed to open index file: %s\n", hmll_strerr(ctx.error)); + return 1; + } + + num_files = hmll_safetensors_index(&ctx, ®istry, index_src); + hmll_source_close(&index_src); + + if (num_files == 0) { + fprintf(stderr, "Failed to parse index: %s\n", hmll_strerr(ctx.error)); + return 2; + } + + char dir[4096]; + get_dir_prefix(path, dir, sizeof(dir)); + + sources = calloc(num_files, sizeof(hmll_source_t)); + if (!sources) { + fprintf(stderr, "Allocation failed\n"); + hmll_free_registry(®istry); + return 2; + } + + for (size_t i = 0; i < num_files; ++i) { + char shard_path[4096 + 64]; + snprintf(shard_path, sizeof(shard_path), "%smodel-%05zu-of-%05zu.safetensors", + dir, i + 1, num_files); + if (hmll_check(hmll_source_open(shard_path, &sources[i]))) { + fprintf(stderr, "Failed to open shard %zu (%s): %s\n", + i + 1, shard_path, hmll_strerr(ctx.error)); + for (size_t j = 0; j < i; ++j) hmll_source_close(&sources[j]); + free(sources); + hmll_free_registry(®istry); + return 2; + } + } + + size_t offset = 0; + for (size_t i = 0; i < num_files; ++i) { + const size_t n = hmll_safetensors_populate_registry(&ctx, ®istry, sources[i], i, offset); + if (n == 0) { + fprintf(stderr, "Failed to populate registry from shard %zu: %s\n", + i + 1, hmll_strerr(ctx.error)); + for (size_t j = 0; j < num_files; ++j) hmll_source_close(&sources[j]); + free(sources); + hmll_free_registry(®istry); + return 2; + } + offset += n; + } + } else { + sources = calloc(1, sizeof(hmll_source_t)); + if (!sources) { + fprintf(stderr, "Allocation failed\n"); + return 2; + } + + if (hmll_check(hmll_source_open(path, &sources[0]))) { + fprintf(stderr, "Failed to open file: %s\n", hmll_strerr(ctx.error)); + free(sources); + return 1; + } + + if (hmll_safetensors_populate_registry(&ctx, ®istry, sources[0], 0, 0) == 0) { + fprintf(stderr, "Failed to populate registry: %s\n", hmll_strerr(ctx.error)); + hmll_source_close(&sources[0]); + free(sources); + return 2; + } + } + + printf("Registry: %zu tensor(s) across %zu file(s)\n", registry.num_tensors, num_files); + +#if defined(__HMLL_CUDA_ENABLED__) + const struct hmll_device device = hmll_device_cuda(0); +#else + const struct hmll_device device = hmll_device_cpu(); +#endif - if (hmll_check(hmll_loader_init(&ctx, &src, 1, hmll_device_cpu(), HMLL_FETCHER_IO_URING))) + if (hmll_check(hmll_loader_init(&ctx, sources, num_files, device, HMLL_FETCHER_IO_URING))) { + fprintf(stderr, "Failed to init loader: %s\n", hmll_strerr(ctx.error)); + for (size_t i = 0; i < num_files; ++i) hmll_source_close(&sources[i]); + free(sources); + hmll_free_registry(®istry); return 3; + } + + timespec_t total_start, total_end; + tick(&total_start); + + size_t total_bytes = 0; + size_t total_errors = 0; + + /* + * Core fetchv strategy: group all tensors that belong to the same shard and + * submit them in a single hmll_fetchv call. This lets the io_uring backend + * pipeline all reads for a shard in one submission burst instead of issuing + * individual SQEs one tensor at a time. + */ + for (size_t fid = 0; fid < num_files; ++fid) { + + /* --- first pass: count tensors for this shard --- */ + size_t shard_n = 0; + for (size_t t = 0; t < registry.num_tensors; ++t) { + if (registry.indexes[t] == (unsigned short)fid) + ++shard_n; + } + if (shard_n == 0) continue; - const hmll_lookup_result_t lookup = hmll_lookup_tensor(&ctx, ®istry, TENSOR_NAME); - if (hmll_success(ctx.error) && lookup.specs != NULL) - { - const hmll_range_t range = (struct hmll_range){ lookup.specs->start, lookup.specs->end }; - const hmll_iobuf_t buffer = hmll_get_buffer_for_range(&ctx, range); - if (hmll_success(ctx.error)) { - // create ranges - const uintptr_t ptr = lookup.specs->start; - struct hmll_range ranges[4] = { - {ptr + 0 * sizeof(float), ptr + 1 * sizeof(float)}, - {ptr + 3 * sizeof(float), ptr + 4 * sizeof(float)}, - {ptr + 2 * sizeof(float), ptr + 3 * sizeof(float)}, - {ptr + 1 * sizeof(float), ptr + 2 * sizeof(float)}, - }; - - struct hmll_iobuf dsts[4] = { - hmll_slice_buffer(&buffer, ranges[0]), - hmll_slice_buffer(&buffer, ranges[1]), - hmll_slice_buffer(&buffer, ranges[2]), - hmll_slice_buffer(&buffer, ranges[3]), - }; - - // Start timing - timespec_t start, end; - tick(&start); - - if (hmll_fetchv(&ctx, lookup.file, dsts, ranges, 4) < 0) { - fprintf(stderr, "Failed to fetch data: %s", hmll_strerr(ctx.error)); - return 4; + /* --- allocate per-shard arrays --- */ + hmll_iobuf_t *dsts = calloc(shard_n, sizeof(hmll_iobuf_t)); + size_t *offsets = calloc(shard_n, sizeof(size_t)); + size_t *tids = calloc(shard_n, sizeof(size_t)); /* original tensor indices */ + + if (!dsts || !offsets || !tids) { + fprintf(stderr, "Shard %zu: allocation failed\n", fid); + free(dsts); free(offsets); free(tids); + total_errors += shard_n; + continue; + } + + /* --- second pass: build buffers and offsets arrays --- */ + size_t k = 0; + size_t shard_bytes = 0; + int alloc_ok = 1; + + for (size_t t = 0; t < registry.num_tensors && alloc_ok; ++t) { + if (registry.indexes[t] != (unsigned short)fid) continue; + + const struct hmll_tensor_specs *specs = ®istry.tensors[t]; + const hmll_range_t range = {specs->start, specs->end}; + + dsts[k] = hmll_get_buffer_for_range(&ctx, range); + if (!hmll_success(ctx.error)) { + fprintf(stderr, "Shard %zu: buffer alloc failed for '%s': %s\n", + fid, registry.names[t], hmll_strerr(ctx.error)); + ctx.error = HMLL_OK; + /* free already-allocated buffers in this shard */ + for (size_t j = 0; j < k; ++j) hmll_free_buffer(&dsts[j]); + alloc_ok = 0; + total_errors += shard_n; + break; } - // End timing and calculate elapsed time - tick(&end); - const double elapsed_ns = time_diff_ns(&start, &end); - const double elapsed_ms = elapsed_ns / 1e6; - const double elapsed_s = elapsed_ns / 1e9; - - if (hmll_success(ctx.error)) { - // Calculate throughput - const double size_mb = (double)(buffer.size) / (1024.0 * 1024.0); - const double throughput_mbps = size_mb / elapsed_s; - - printf("Fetch completed in %.3f ms (%.6f s)\n", elapsed_ms, elapsed_s); - printf("Tensor size: %.2f MB\n", size_mb); - printf("Throughput: %.2f MB/s\n", throughput_mbps); - - unsigned short *bf16_ptr; - if (hmll_device_is_cuda(ctx.fetcher->device)) { - bf16_ptr = malloc(buffer.size); - cudaMemcpy(bf16_ptr, buffer.ptr, hmll_numel(lookup.specs) * sizeof(short), cudaMemcpyDeviceToHost); - } else { - bf16_ptr = buffer.ptr; - } - - unsigned long sum = 0; - for (size_t i = 0; i < hmll_numel(lookup.specs); ++i) sum += bf16_ptr[i]; - - printf("Sum: %lu\n", sum); - } else { - printf("Got an error while reading the safetensors: %s\n", hmll_strerr(ctx.error)); + offsets[k] = specs->start; + tids[k] = t; + shard_bytes += specs->end - specs->start; + ++k; + } + + if (!alloc_ok) { + free(dsts); free(offsets); free(tids); + continue; + } + + /* --- single fetchv call for the whole shard --- */ + printf("\n[Shard %zu/%zu] fetching %zu tensor(s) (%.2f MB) in one call\n", + fid + 1, num_files, shard_n, + (double)shard_bytes / (1024.0 * 1024.0)); + + timespec_t shard_start, shard_end; + tick(&shard_start); + + const ssize_t fetched = hmll_fetchv(&ctx, (int)fid, dsts, offsets, shard_n); + + tick(&shard_end); + + if (fetched < 0 || !hmll_success(ctx.error)) { + fprintf(stderr, "Shard %zu: fetchv failed: %s\n", + fid, hmll_strerr(ctx.error)); + ctx.error = HMLL_OK; + total_errors += shard_n; + } else { + const double shard_elapsed_s = time_diff_ns(&shard_start, &shard_end) / 1e9; + const double shard_mb = (double)shard_bytes / (1024.0 * 1024.0); + + total_bytes += (size_t)fetched; + + /* per-tensor report */ + for (size_t j = 0; j < shard_n; ++j) { + const size_t t = tids[j]; + const double mb = (double)dsts[j].size / (1024.0 * 1024.0); + printf(" [%zu/%zu] %-60s %8.2f MB\n", + t + 1, registry.num_tensors, registry.names[t], mb); } + + printf("Shard %zu throughput: %.2f MB/s (%.3f s)\n", + fid + 1, shard_mb / shard_elapsed_s, shard_elapsed_s); } + + for (size_t j = 0; j < shard_n; ++j) hmll_free_buffer(&dsts[j]); + free(dsts); free(offsets); free(tids); } - return 0; + tick(&total_end); + + const double total_elapsed_s = time_diff_ns(&total_start, &total_end) / 1e9; + const double total_mb = (double)total_bytes / (1024.0 * 1024.0); + + printf("\n=== Summary ===\n"); + printf("Tensors fetched : %zu (errors: %zu)\n", + registry.num_tensors - total_errors, total_errors); + printf("Total data : %.2f MB\n", total_mb); + printf("Total time : %.3f s\n", total_elapsed_s); + printf("Throughput : %.2f MB/s\n", total_mb / total_elapsed_s); + + hmll_free_registry(®istry); + for (size_t i = 0; i < num_files; ++i) hmll_source_close(&sources[i]); + free(sources); + + return total_errors > 0 ? 5 : 0; } diff --git a/examples/safetensors.c b/examples/safetensors.c index 01e7f99..bd2bb7e 100644 --- a/examples/safetensors.c +++ b/examples/safetensors.c @@ -1,5 +1,6 @@ #include #include +#include #include #ifdef _WIN32 @@ -32,83 +33,190 @@ static double time_diff_ns(const timespec_t *start, const timespec_t *end) { #include #endif -#define TENSOR_NAME "language_model.model.embed_tokens.weight" -// #define TENSOR_NAME "model.embed_tokens.weight" +static int path_ends_with(const char *path, const char *suffix) { + const size_t plen = strlen(path); + const size_t slen = strlen(suffix); + return plen >= slen && strcmp(path + plen - slen, suffix) == 0; +} + +static void get_dir_prefix(const char *path, char *dir, size_t dir_size) { + strncpy(dir, path, dir_size - 1); + dir[dir_size - 1] = '\0'; + char *last_slash = strrchr(dir, '/'); + if (last_slash) { + *(last_slash + 1) = '\0'; + } else { + dir[0] = '\0'; + } +} int main(const int argc, const char** argv) { if (argc < 2) { - printf("No file specified.\nInvoke through hmll_safetensors_ex "); + fprintf(stderr, "Usage: hmll_safetensors_ex \n"); return 1; } - hmll_t ctx = {0}; - hmll_source_t src = {0}; - if (hmll_check(hmll_source_open(argv[1], &src))) - return 1; + const char *path = argv[1]; + const int is_sharded = path_ends_with(path, ".index.json"); - // Read safetensors table with all tensors mapping + hmll_t ctx = {0}; hmll_registry_t registry = {0}; - if (hmll_safetensors_populate_registry(&ctx, ®istry, src, 0, 0) == 0) - return 2; + hmll_source_t *sources = NULL; + size_t num_files = 1; + + if (is_sharded) { + hmll_source_t index_src = {0}; + if (hmll_check(hmll_source_open(path, &index_src))) { + fprintf(stderr, "Failed to open index file: %s\n", hmll_strerr(ctx.error)); + return 1; + } - if (hmll_check(hmll_loader_init(&ctx, &src, 1, hmll_device_cuda(0), HMLL_FETCHER_IO_URING))) - return 3; + num_files = hmll_safetensors_index(&ctx, ®istry, index_src); + hmll_source_close(&index_src); + + if (num_files == 0) { + fprintf(stderr, "Failed to parse index: %s\n", hmll_strerr(ctx.error)); + return 2; + } - const hmll_lookup_result_t lookup = hmll_lookup_tensor(&ctx, ®istry, TENSOR_NAME); - if (hmll_success(ctx.error) && lookup.specs) - { - const hmll_range_t range = (struct hmll_range){ lookup.specs->start, lookup.specs->end }; - const hmll_iobuf_t buffer = hmll_get_buffer_for_range(&ctx, range); - if (hmll_success(ctx.error)) { - // Start timing - timespec_t start, end; - tick(&start); - - if (hmll_fetch(&ctx, lookup.file, &buffer, range.start) < range.end - range.start) { - fprintf(stderr, "Failed to fetch data: %s", hmll_strerr(ctx.error)); - return 4; + char dir[4096]; + get_dir_prefix(path, dir, sizeof(dir)); + + sources = calloc(num_files, sizeof(hmll_source_t)); + if (!sources) { + fprintf(stderr, "Allocation failed\n"); + hmll_free_registry(®istry); + return 2; + } + + for (size_t i = 0; i < num_files; ++i) { + char shard_path[4096 + 64]; + snprintf(shard_path, sizeof(shard_path), "%smodel-%05zu-of-%05zu.safetensors", + dir, i + 1, num_files); + if (hmll_check(hmll_source_open(shard_path, &sources[i]))) { + fprintf(stderr, "Failed to open shard %zu (%s): %s\n", + i + 1, shard_path, hmll_strerr(ctx.error)); + for (size_t j = 0; j < i; ++j) hmll_source_close(&sources[j]); + free(sources); + hmll_free_registry(®istry); + return 2; } + } - // End timing and calculate elapsed time - tick(&end); - const double elapsed_ns = time_diff_ns(&start, &end); - const double elapsed_ms = elapsed_ns / 1e6; - const double elapsed_s = elapsed_ns / 1e9; - - if (hmll_success(ctx.error)) { - // Calculate throughput - const double size_mb = (double)(buffer.size) / (1024.0 * 1024.0); - const double throughput_mbps = size_mb / elapsed_s; - - printf("Fetch completed in %.3f ms (%.6f s)\n", elapsed_ms, elapsed_s); - printf("Tensor size: %.2f MB\n", size_mb); - printf("Throughput: %.2f MB/s\n", throughput_mbps); - - unsigned short *bf16_ptr; - if (hmll_device_is_cuda(ctx.fetcher->device)) { - bf16_ptr = malloc(buffer.size); - cudaMemcpy(bf16_ptr, buffer.ptr, hmll_numel(lookup.specs) * sizeof(short), cudaMemcpyDeviceToHost); - } else { - bf16_ptr = buffer.ptr; - } - - unsigned long sum = 0; - for (size_t i = 0; i < hmll_numel(lookup.specs); ++i) sum += bf16_ptr[i]; - - printf("Sum: %lu\n", sum); - } else { - printf("Got an error while reading the safetensors: %s\n", hmll_strerr(ctx.error)); + size_t offset = 0; + for (size_t i = 0; i < num_files; ++i) { + const size_t n = hmll_safetensors_populate_registry(&ctx, ®istry, sources[i], i, offset); + if (n == 0) { + fprintf(stderr, "Failed to populate registry from shard %zu: %s\n", + i + 1, hmll_strerr(ctx.error)); + for (size_t j = 0; j < num_files; ++j) hmll_source_close(&sources[j]); + free(sources); + hmll_free_registry(®istry); + return 2; } + offset += n; } } else { - if (!lookup.specs) - fprintf(stderr, "Tensor not found in safetensors file.\n"); - else - fprintf(stderr, "Failed to lookup tensor: %s\n", hmll_strerr(ctx.error)); + sources = calloc(1, sizeof(hmll_source_t)); + if (!sources) { + fprintf(stderr, "Allocation failed\n"); + return 2; + } + + if (hmll_check(hmll_source_open(path, &sources[0]))) { + fprintf(stderr, "Failed to open file: %s\n", hmll_strerr(ctx.error)); + free(sources); + return 1; + } + + if (hmll_safetensors_populate_registry(&ctx, ®istry, sources[0], 0, 0) == 0) { + fprintf(stderr, "Failed to populate registry: %s\n", hmll_strerr(ctx.error)); + hmll_source_close(&sources[0]); + free(sources); + return 2; + } } + printf("Registry: %zu tensor(s) across %zu file(s)\n", registry.num_tensors, num_files); + +#if defined(__HMLL_CUDA_ENABLED__) + const struct hmll_device device = hmll_device_cuda(0); +#else + const struct hmll_device device = hmll_device_cpu(); +#endif + + if (hmll_check(hmll_loader_init(&ctx, sources, num_files, device, HMLL_FETCHER_IO_URING))) { + fprintf(stderr, "Failed to init loader: %s\n", hmll_strerr(ctx.error)); + for (size_t i = 0; i < num_files; ++i) hmll_source_close(&sources[i]); + free(sources); + hmll_free_registry(®istry); + return 3; + } + + timespec_t total_start, total_end; + tick(&total_start); + + size_t total_bytes = 0; + size_t errors = 0; + + for (size_t t = 0; t < registry.num_tensors; ++t) { + const char *name = registry.names[t]; + const struct hmll_tensor_specs *specs = ®istry.tensors[t]; + const unsigned short file_idx = registry.indexes[t]; + const hmll_range_t range = {specs->start, specs->end}; + + hmll_iobuf_t buffer = hmll_get_buffer_for_range(&ctx, range); + if (!hmll_success(ctx.error)) { + fprintf(stderr, "[%zu/%zu] %s: buffer alloc failed: %s\n", + t + 1, registry.num_tensors, name, hmll_strerr(ctx.error)); + ctx.error = HMLL_OK; + ++errors; + continue; + } + + timespec_t t_start, t_end; + tick(&t_start); + + const ssize_t fetched = hmll_fetch(&ctx, file_idx, &buffer, range.start); + + tick(&t_end); + + if (fetched < 0 || !hmll_success(ctx.error)) { + fprintf(stderr, "[%zu/%zu] %s: fetch failed: %s\n", + t + 1, registry.num_tensors, name, hmll_strerr(ctx.error)); + hmll_free_buffer(&buffer); + ctx.error = HMLL_OK; + ++errors; + continue; + } + + const double elapsed_s = time_diff_ns(&t_start, &t_end) / 1e9; + const double size_mb = (double)(range.end - range.start) / (1024.0 * 1024.0); + const double throughput_mbps = size_mb / elapsed_s; + + total_bytes += (size_t)fetched; + + printf("[%zu/%zu] %-60s %8.2f MB %9.2f MB/s\n", + t + 1, registry.num_tensors, name, size_mb, throughput_mbps); + + hmll_free_buffer(&buffer); + } + + tick(&total_end); + + const double total_elapsed_s = time_diff_ns(&total_start, &total_end) / 1e9; + const double total_mb = (double)total_bytes / (1024.0 * 1024.0); + + printf("\n=== Summary ===\n"); + printf("Tensors fetched : %zu (errors: %zu)\n", registry.num_tensors - errors, errors); + printf("Total data : %.2f MB\n", total_mb); + printf("Total time : %.3f s\n", total_elapsed_s); + printf("Throughput : %.2f MB/s\n", total_mb / total_elapsed_s); + hmll_free_registry(®istry); - hmll_source_close(&src); - return 0; + for (size_t i = 0; i < num_files; ++i) hmll_source_close(&sources[i]); + free(sources); + + return errors > 0 ? 5 : 0; } diff --git a/include/hmll/linux/backend/iouring.h b/include/hmll/linux/backend/iouring.h index 49eff45..828ebc2 100644 --- a/include/hmll/linux/backend/iouring.h +++ b/include/hmll/linux/backend/iouring.h @@ -11,9 +11,12 @@ #define HMLL_URING_BUFFER_SIZE (512U * 1024) #endif +#include #include #include "hmll/types.h" +static_assert(HMLL_URING_BUFFER_SIZE % 4096 == 0, "HMLL_URING_BUFFER_SIZE should be 4096-aligned"); + struct hmll_iouring_iobusy { unsigned long long bits[HMLL_URING_IOBUSY_WORDS]; @@ -88,7 +91,6 @@ struct hmll_io_uring { struct io_uring ioring; struct iovec *iovecs; struct hmll_iouring_iobusy iobusy; - struct hmll_iouring_cca iocca; // congestion control // store optional device data void *device_ctx; @@ -119,6 +121,18 @@ static inline void hmll_io_uring_slot_set_available(struct hmll_iouring_iobusy * iobusy->bits[slot >> 6] &= ~(1ULL << (slot & 63)); } +/** + * Map a user-visible source index to the io_uring registered-file index + * for the buffered (page-cache) fd. Layout: [b0, d0, b1, d1, ...]. + */ +static inline unsigned hmll_io_uring_buffered_fd(const unsigned iofile) { return iofile * 2; } + +/** + * Map a user-visible source index to the io_uring registered-file index + * for the O_DIRECT fd. + */ +static inline unsigned hmll_io_uring_direct_fd(const unsigned iofile) { return iofile * 2 + 1; } + struct hmll_error hmll_io_uring_init(struct hmll *, struct hmll_device); void hmll_io_uring_destroy(void *backend); #endif // HMLL_FETCHER_IOURING_H diff --git a/include/hmll/loader.h b/include/hmll/loader.h index 617c0b3..d89f2d4 100644 --- a/include/hmll/loader.h +++ b/include/hmll/loader.h @@ -19,8 +19,8 @@ struct hmll_loader struct hmll_device device; void *backend_impl_; void(*backend_free)(void *backend); - ssize_t(*fetch_range_impl_)(struct hmll *, int, const struct hmll_iobuf *, size_t); - ssize_t(*fetchv_range_impl_)(struct hmll *, int, const struct hmll_iobuf *, const size_t *, size_t); + ssize_t(*fetch_range_impl_)(struct hmll *, unsigned, const struct hmll_iobuf *, size_t); + ssize_t(*fetchv_range_impl_)(struct hmll *, unsigned, const struct hmll_iobuf *, const size_t *, size_t); }; typedef struct hmll_loader hmll_loader_t; diff --git a/include/hmll/unix/file.h b/include/hmll/unix/file.h index af51417..0f9f99c 100644 --- a/include/hmll/unix/file.h +++ b/include/hmll/unix/file.h @@ -6,6 +6,7 @@ struct hmll_source { int fd; + int d_fd; /* O_DIRECT fd for aligned I/O (Linux only, -1 elsewhere) */ size_t size; const unsigned char *content; }; diff --git a/lib/linux/backend/iouring.c b/lib/linux/backend/iouring.c index 5100be4..ef2fb95 100644 --- a/lib/linux/backend/iouring.c +++ b/lib/linux/backend/iouring.c @@ -1,10 +1,12 @@ #include +#include #include "hmll/hmll.h" #include "hmll/memory.h" #include "hmll/linux/backend/iouring.h" #include "sys/mman.h" #define HMLL_IO_URING_FADVISE_TAG (1ULL << 63) +#define HMLL_IO_URING_DEFAULT_QUEUE_PARAMS IORING_SETUP_SQPOLL #if defined(__HMLL_CUDA_ENABLED__) #include "hmll/cuda.h" @@ -12,8 +14,59 @@ #include #endif +/* ── runtime kernel version detection ───────────────────────────────── */ +static inline unsigned hmll_kernel_version_internal(unsigned maj, unsigned min) +{ + return (maj << 16) | min; +} + +static unsigned hmll_kernel_version(void) +{ + static unsigned cached = 0; + if (cached) return cached; + + struct utsname u; + if (uname(&u) != 0) return 0; + + unsigned maj = 0, min = 0; + if (sscanf(u.release, "%u.%u", &maj, &min) < 2) return 0; + + cached = hmll_kernel_version_internal(maj, min); + return cached; +} -static inline int hmll_io_uring_get_setup_flags(void) { return IORING_SETUP_SQPOLL; } +/** + * Build io_uring setup flags based on the running kernel version. + * + * SQPOLL (always) — kernel thread polls SQ, eliminates submit syscalls + * COOP_TASKRUN (>= 6.1) — no IPI for CQE delivery, process on next syscall + * TASKRUN_FLAG (>= 6.1) — companion: sets SQ flag when CQEs are pending + */ +static inline int hmll_io_uring_get_setup_flags(void) +{ + const unsigned kversion = hmll_kernel_version(); + + int flags = HMLL_IO_URING_DEFAULT_QUEUE_PARAMS; + if (kversion >= hmll_kernel_version_internal(6, 1)) { + flags |= IORING_SETUP_COOP_TASKRUN | IORING_SETUP_TASKRUN_FLAG; + } + + return flags; +} + +/* ── scratch allocator (stack with heap fallback) ───────────────────── */ + +static inline void *hmll_scratch_alloc( + uint8_t *stack, const size_t size, const size_t need, void **to_free +) { + *to_free = NULL; + if (need <= size) return stack; + void *p = calloc(1, need); + *to_free = p; + return p; +} + +/* ── staging buffer registration ────────────────────────────────────── */ static struct hmll_error hmll_io_uring_register_staging_buffers( struct hmll *ctx, @@ -46,11 +99,24 @@ static struct hmll_error hmll_io_uring_register_staging_buffers( return HMLL_OK; } +/* ── fadvise helper ─────────────────────────────────────────────────── */ + +static inline void hmll_io_uring_queue_fadvise( + struct hmll_io_uring *fetcher, const unsigned iofd, const size_t off, const size_t len +) { + struct io_uring_sqe *sqe = io_uring_get_sqe(&fetcher->ioring); + if (!sqe) return; + io_uring_prep_fadvise(sqe, (int)iofd, off, len, POSIX_FADV_WILLNEED); + io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); + io_uring_sqe_set_data64(sqe, HMLL_IO_URING_FADVISE_TAG); +} + +/* ── CUDA helpers ───────────────────────────────────────────────────── */ + static inline void hmll_io_uring_sync(const struct hmll_device device, const struct hmll_io_uring *fetcher) { if (hmll_device_is_cuda(device)) { #ifdef __HMLL_CUDA_ENABLED__ - // Wait for all pending CUDA operations to complete for (size_t i = 0; i < HMLL_URING_QUEUE_DEPTH; ++i) { if (hmll_io_uring_slot_is_busy(fetcher->iobusy, i)) { const struct hmll_io_uring_cuda_context *cd = (struct hmll_io_uring_cuda_context *)fetcher->device_ctx + i; @@ -60,14 +126,9 @@ static inline void hmll_io_uring_sync(const struct hmll_device device, const str } #endif } - HMLL_UNUSED(fetcher); } -/** - * Checks for completed CUDA events and reclaims the associated io_uring slots. - * If CUDA is disabled or the device is CPU, this is a no-op. - */ static inline void hmll_io_uring_reclaim_slots( struct hmll_io_uring *fetcher, const struct hmll_device device @@ -76,8 +137,6 @@ static inline void hmll_io_uring_reclaim_slots( if (!hmll_device_is_cuda(device)) return; struct hmll_io_uring_cuda_context *dctx = fetcher->device_ctx; - - // TODO(mfuntowicz): Should we directly store `slots` which are doing memcpy currently to avoid full scan? 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)) { @@ -94,9 +153,32 @@ static inline void hmll_io_uring_reclaim_slots( } /** - * Prepares a single SQE (Submission Queue Entry). - * Handles the difference between direct CPU buffer reads and CUDA staging buffer reads. + * When all slots are busy with CUDA memcpy, synchronously wait on one event + * to free a slot. No-op when a slot is already available. */ +static inline void hmll_io_uring_cuda_relieve_pressure(struct hmll_io_uring *fetcher) +{ +#ifdef __HMLL_CUDA_ENABLED__ + if (hmll_io_uring_slot_find_available(fetcher->iobusy) >= 0) return; + + struct hmll_io_uring_cuda_context *dctx = fetcher->device_ctx; + for (size_t i = 0; i < HMLL_URING_QUEUE_DEPTH; i++) { + if (!hmll_io_uring_slot_is_busy(fetcher->iobusy, i)) continue; + 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); + return; + } + } +#else + HMLL_UNUSED(fetcher); +#endif +} + +/* ── SQE / CQE primitives ──────────────────────────────────────────── */ + static inline void hmll_io_uring_prep_sqe( const struct hmll_io_uring *fetcher, const struct hmll_device device, @@ -104,28 +186,23 @@ static inline void hmll_io_uring_prep_sqe( void *dst, const size_t offset, const size_t len, - const unsigned short iofile, + const int iofile, const int slot ) { if (hmll_device_is_cpu(device)) { - // CPU: Read directly into user memory io_uring_prep_read(sqe, iofile, dst, len, offset); io_uring_sqe_set_data64(sqe, slot); } #if defined(__HMLL_CUDA_ENABLED__) else if (hmll_device_is_cuda(device)) { - // CUDA: Read into registered staging buffers struct hmll_io_uring_cuda_context *dctx = fetcher->device_ctx; - void *buf = fetcher->iovecs[slot].iov_base; - dctx[slot].offset = offset; - io_uring_prep_read_fixed(sqe, iofile, buf, len, offset, slot); + io_uring_prep_read_fixed(sqe, iofile, fetcher->iovecs[slot].iov_base, len, offset, slot); io_uring_sqe_set_data(sqe, dctx + slot); } #else HMLL_UNUSED(fetcher); #endif - io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); } @@ -141,11 +218,6 @@ static inline int hmll_io_uring_get_sqe(struct hmll_io_uring *fetcher, struct io return slot; } -/** - * Handles the completion of an IO request (CQE). - * For CPU: just marks a slot available. - * For CUDA: Dispatches the Async Memcpy from staging to GPU. - */ static inline void hmll_io_uring_handle_completion( struct hmll_io_uring *fetcher, const struct io_uring_cqe *cqe, @@ -154,16 +226,13 @@ static inline void hmll_io_uring_handle_completion( const int32_t len ) { if (hmll_device_is_cpu(dst->device)) { - const uint64_t cb_slot = cqe->user_data; - hmll_io_uring_slot_set_available(&fetcher->iobusy, cb_slot); + hmll_io_uring_slot_set_available(&fetcher->iobusy, cqe->user_data); } #if defined(__HMLL_CUDA_ENABLED__) else if (hmll_device_is_cuda(dst->device)) { struct hmll_io_uring_cuda_context *cctx = (struct hmll_io_uring_cuda_context *)cqe->user_data; - - void *to = (char *)dst->ptr + (cctx->offset - offset); + void *to = (char *)dst->ptr + (cctx->offset - offset); void *from = fetcher->iovecs[cctx->slot].iov_base; - cudaMemcpyAsync(to, from, len, cudaMemcpyHostToDevice, cctx->stream); cudaEventRecord(cctx->done, cctx->stream); hmll_io_uring_cuda_stream_set_memcpy(&cctx->state); @@ -174,33 +243,35 @@ static inline void hmll_io_uring_handle_completion( #endif } -static ssize_t hmll_io_uring_fetch_impl( +/** + * Generic fetch loop for a single buffer. When @p fadvise is non-zero an + * initial POSIX_FADV_WILLNEED is queued; suppress it (pass 0) when the caller + * has already arranged readahead or when using an O_DIRECT fd. + */ +static ssize_t hmll_io_uring_fetch_loop( struct hmll *ctx, - const int iofile, + const int iofd, const struct hmll_iobuf *dst, - const size_t offset + const size_t offset, + const int fadvise ) { - if (hmll_check(ctx->error)) return -1; - struct hmll_io_uring *fetcher = ctx->fetcher->backend_impl_; - size_t n_dma = 0; + size_t n_inflight = 0; size_t b_read = 0; size_t b_submitted = 0; struct io_uring_cqe *cqes[HMLL_URING_QUEUE_DEPTH]; - 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_WILLNEED); - io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); - io_uring_sqe_set_data64(sqe, HMLL_IO_URING_FADVISE_TAG); - } + if (fadvise) + hmll_io_uring_queue_fadvise(fetcher, iofd, offset, dst->size); while (b_read < dst->size) { hmll_io_uring_reclaim_slots(fetcher, dst->device); - while (b_submitted < dst->size) { + /* ── submit ── */ + while (b_submitted < dst->size && n_inflight < HMLL_URING_QUEUE_DEPTH) { + struct io_uring_sqe *sqe = NULL; + int slot; if (unlikely((slot = hmll_io_uring_get_sqe(fetcher, &sqe)) < 0)) break; @@ -208,144 +279,257 @@ static ssize_t hmll_io_uring_fetch_impl( const size_t to_read = (remaining < HMLL_URING_BUFFER_SIZE) ? remaining : HMLL_URING_BUFFER_SIZE; const size_t file_offset = offset + b_submitted; - hmll_io_uring_prep_sqe(fetcher, dst->device, sqe, (char *)dst->ptr + b_submitted, file_offset, to_read, iofile, slot); + hmll_io_uring_prep_sqe(fetcher, dst->device, sqe, (char *)dst->ptr + b_submitted, file_offset, to_read, iofd, slot); b_submitted += to_read; - ++n_dma; + n_inflight++; } - // update congestion control algorithm - if (likely(n_dma > 0)) { - const size_t nwait = n_dma < fetcher->iocca.window ? n_dma : fetcher->iocca.window; + io_uring_submit(&fetcher->ioring); - struct timespec ts_start, ts_end; - clock_gettime(CLOCK_MONOTONIC, &ts_start); + /* if slots are exhausted by CUDA memcpy, relieve pressure before blocking */ + if (b_submitted < dst->size) + hmll_io_uring_cuda_relieve_pressure(fetcher); - if (unlikely(io_uring_submit_and_wait(&fetcher->ioring, nwait) < 0)) { - // todo: do we need to reset the cca? hmll_io_uring_cca_init(&fetcher->iocca) + /* ── complete: non-blocking peek first, block only when pipeline full ── */ + unsigned count = io_uring_peek_batch_cqe(&fetcher->ioring, cqes, HMLL_URING_QUEUE_DEPTH); + if (count == 0 && n_inflight > 0) { + struct io_uring_cqe *cqe; + if (unlikely(io_uring_wait_cqe(&fetcher->ioring, &cqe) < 0)) { ctx->error = HMLL_ERR(HMLL_ERR_IO_ERROR); return -1; } - clock_gettime(CLOCK_MONOTONIC, &ts_end); - - // todo: approximated version of the number of bytes actually reads because it assumes full reads - hmll_io_uring_cca_update(&fetcher->iocca, HMLL_URING_BUFFER_SIZE * nwait, ts_start, ts_end); + cqes[0] = cqe; + count = 1; } - unsigned count = 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]; - if (unlikely(cqe->user_data == HMLL_IO_URING_FADVISE_TAG)) - continue; - - --n_dma; - if (unlikely(cqe->res < 0)) { - ctx->error = HMLL_SYS_ERR(-cqe->res); - io_uring_cq_advance(&fetcher->ioring, count); - return -1; - } + for (unsigned i = 0; i < count; i++) { + const struct io_uring_cqe *cqe = cqes[i]; + if (unlikely(cqe->user_data == HMLL_IO_URING_FADVISE_TAG)) + continue; - b_read += cqe->res; - hmll_io_uring_handle_completion(fetcher, cqe, dst, offset, cqe->res); + if (unlikely(cqe->res < 0)) { + ctx->error = HMLL_SYS_ERR(-cqe->res); + io_uring_cq_advance(&fetcher->ioring, count); + return -1; } - io_uring_cq_advance(&fetcher->ioring, count); + b_read += cqe->res; + n_inflight--; + hmll_io_uring_handle_completion(fetcher, cqe, dst, offset, cqe->res); } + io_uring_cq_advance(&fetcher->ioring, count); } hmll_io_uring_sync(dst->device, fetcher); return (ssize_t)b_read; } -static ssize_t hmll_io_uring_fetchv_impl( +#if defined(__HMLL_CUDA_ENABLED__) +/** + * CUDA split-I/O fetch: O_DIRECT for the page-aligned core, buffered I/O for + * the unaligned head/tail. fadvise(WILLNEED) is issued for the edge regions + * before reading the core, giving the kernel time to readahead. + */ +static ssize_t hmll_io_uring_fetch_cuda_split( struct hmll *ctx, - const int iofile, + const int bfd, + const int dfd, + const struct hmll_iobuf *dst, + const size_t offset +) { + const size_t end = offset + dst->size; + const size_t aligned_start = ALIGN_UP(offset, ALIGN_PAGE); + const size_t aligned_end = ALIGN_DOWN(end, ALIGN_PAGE); + + const size_t head_size = aligned_start - offset; + const size_t tail_size = end - aligned_end; + const size_t core_size = (aligned_end > aligned_start) ? aligned_end - aligned_start : 0; + + if (core_size < ALIGN_PAGE * 2) + return hmll_io_uring_fetch_loop(ctx, bfd, dst, offset, 1); + + struct hmll_io_uring *fetcher = ctx->fetcher->backend_impl_; + ssize_t total = 0; + + if (head_size > 0) hmll_io_uring_queue_fadvise(fetcher, bfd, offset, head_size); + if (tail_size > 0) hmll_io_uring_queue_fadvise(fetcher, bfd, aligned_end, tail_size); + + /* aligned core via O_DIRECT → staging → GPU */ + struct hmll_iobuf core_dst = { + .size = core_size, .ptr = (char *)dst->ptr + head_size, .device = dst->device, + }; + ssize_t n = hmll_io_uring_fetch_loop(ctx, dfd, &core_dst, aligned_start, 0); + if (n < 0) return -1; + total += n; + + /* head edge from buffered fd (fadvise already queued) */ + if (head_size > 0) { + struct hmll_iobuf head_dst = { + .size = head_size, .ptr = dst->ptr, .device = dst->device, + }; + n = hmll_io_uring_fetch_loop(ctx, bfd, &head_dst, offset, 0); + if (n < 0) return -1; + total += n; + } + + /* tail edge from buffered fd (fadvise already queued) */ + if (tail_size > 0) { + struct hmll_iobuf tail_dst = { + .size = tail_size, .ptr = (char *)dst->ptr + head_size + core_size, .device = dst->device, + }; + n = hmll_io_uring_fetch_loop(ctx, bfd, &tail_dst, aligned_end, 0); + if (n < 0) return -1; + total += n; + } + + return total; +} +#endif /* __HMLL_CUDA_ENABLED__ */ + +static ssize_t hmll_io_uring_fetch_impl( + struct hmll *ctx, + const unsigned iofile, + const struct hmll_iobuf *dst, + const size_t offset +) { + if (hmll_check(ctx->error)) return -1; + + const unsigned bfd = hmll_io_uring_buffered_fd(iofile); + +#if defined(__HMLL_CUDA_ENABLED__) + if (hmll_device_is_cuda(dst->device)) { + const int dfd = hmll_io_uring_direct_fd(iofile); + return hmll_io_uring_fetch_cuda_split(ctx, bfd, dfd, dst, offset); + } +#endif + + return hmll_io_uring_fetch_loop(ctx, bfd, dst, offset, 1); +} + +/* ── fetchv CQE handler ─────────────────────────────────────────────── */ + +/* user_data encoding for fetchv CQEs: (bidx << 8) | slot */ +#define FETCHV_BIDX_SHIFT 8 +#define FETCHV_SLOT_MASK ((uint64_t)HMLL_URING_QUEUE_DEPTH - 1) + +/** + * Handle a single fetchv CQE. Frees the slot (CPU) or kicks off the + * staging→GPU memcpy (CUDA). Returns bytes read, or -1 on I/O error. + */ +static inline ssize_t hmll_io_uring_fetchv_handle_cqe( + struct hmll *ctx, + struct hmll_io_uring *fetcher, + const struct io_uring_cqe *cqe, + const struct hmll_iobuf *dsts, + const size_t *slot_offsets, + const int is_cuda +) { + if (cqe->res < 0) { + ctx->error = HMLL_SYS_ERR(-cqe->res); + return -1; + } + + const uint64_t data = cqe->user_data; + const uint32_t slot = (uint32_t)(data & FETCHV_SLOT_MASK); + const uint32_t bidx = (uint32_t)(data >> FETCHV_BIDX_SHIFT); + + if (!is_cuda) { + HMLL_UNUSED(dsts); + HMLL_UNUSED(slot_offsets); + 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]; + void *to = (char *)dsts[bidx].ptr + slot_offsets[slot]; + void *from = fetcher->iovecs[slot].iov_base; + cudaMemcpyAsync(to, from, (size_t)cqe->res, cudaMemcpyHostToDevice, cctx->stream); + cudaEventRecord(cctx->done, cctx->stream); + hmll_io_uring_cuda_stream_set_memcpy(&cctx->state); + } +#else + (void)bidx; +#endif + + return (ssize_t)cqe->res; +} + +/* ── fetchv loop ────────────────────────────────────────────────────── */ + +struct fetchv_buf_state { + size_t submitted; + size_t size; + unsigned char fadvise_sent; +}; + +/** + * Generic fetchv loop: reads all buffers through a single registered file index. + * Used for the CPU path and as a building block for the CUDA split-I/O path. + * When @p fadvise is 0, per-buffer POSIX_FADV_WILLNEED is suppressed (useful + * when the caller already issued readahead or when reading via O_DIRECT). + */ +static ssize_t hmll_io_uring_fetchv_loop( + struct hmll *ctx, + const unsigned iofd, const struct hmll_iobuf *dsts, const size_t *offsets, - const size_t n + const size_t n, + const int fadvise ) { if (hmll_check(ctx->error)) return -1; 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); + const int is_cuda = hmll_device_is_cuda(dsts[0].device); - /* user_data encoding for CQEs: high bit = fadvise (skip), else (bidx << 8) | slot */ - 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; - }; - - /* Scratch layout: [buf_states][active_indices][slot_offsets] */ + /* scratch: [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; - - 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 { - void *p = calloc(1, scratch_size); - if (!p) { - ctx->error = HMLL_ERR(HMLL_ERR_ALLOCATION_FAILED); - return -1; - } - 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; + void *scratch_to_free; + uint8_t *scratch = hmll_scratch_alloc( + stack_scratch, sizeof(stack_scratch), sz_state + sz_idx + sz_slot, &scratch_to_free); + if (!scratch) { + ctx->error = HMLL_ERR(HMLL_ERR_ALLOCATION_FAILED); + return -1; } - /* Build list of buffers that have bytes to read */ + struct fetchv_buf_state *buf_states = (struct fetchv_buf_state *)scratch; + uint32_t *active_indices = (uint32_t *)(scratch + sz_state); + size_t *slot_offsets = (size_t *)(scratch + sz_state + sz_idx); + size_t n_active = 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; + buf_states[i] = (struct fetchv_buf_state){ .size = dsts[i].size }; if (dsts[i].size > 0) active_indices[n_active++] = (uint32_t)i; } - 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) { hmll_io_uring_reclaim_slots(fetcher, dsts[0].device); - /* Submit: round-robin over active buffers, send fadvise then chunked reads */ - while (n_active > 0) { + /* ── submit: round-robin across buffers ── */ + while (n_active > 0 && n_in_flight < HMLL_URING_QUEUE_DEPTH) { if (active_cursor >= n_active) active_cursor = 0; const uint32_t bidx = active_indices[active_cursor]; 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, HMLL_IO_URING_FADVISE_TAG); st->fadvise_sent = 1; + if (fadvise) + hmll_io_uring_queue_fadvise(fetcher, iofd, offsets[bidx], st->size); continue; } const int slot = hmll_io_uring_slot_find_available(fetcher->iobusy); if (slot < 0) break; + struct io_uring_sqe *sqe = io_uring_get_sqe(&fetcher->ioring); if (!sqe) break; @@ -355,15 +539,15 @@ static ssize_t hmll_io_uring_fetchv_impl( 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; #if defined(__HMLL_CUDA_ENABLED__) if (is_cuda) - io_uring_prep_read_fixed(sqe, iofile, fetcher->iovecs[slot].iov_base, to_read, file_off, slot); + io_uring_prep_read_fixed(sqe, iofd, fetcher->iovecs[slot].iov_base, to_read, file_off, slot); else - io_uring_prep_read(sqe, iofile, read_dst, to_read, file_off); + io_uring_prep_read(sqe, iofd, (char *)dsts[bidx].ptr + st->submitted, to_read, file_off); + #else - io_uring_prep_read(sqe, iofile, read_dst, to_read, file_off); + io_uring_prep_read(sqe, (int)iofd, (char *)dsts[bidx].ptr + st->submitted, to_read, file_off); #endif io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE); io_uring_sqe_set_data64(sqe, ((uint64_t)bidx << FETCHV_BIDX_SHIFT) | (uint64_t)slot); @@ -371,83 +555,50 @@ static ssize_t hmll_io_uring_fetchv_impl( st->submitted += to_read; n_in_flight++; - if (st->submitted >= st->size) { - n_active--; - active_indices[active_cursor] = active_indices[n_active]; - } else { + if (st->submitted >= st->size) + active_indices[active_cursor] = active_indices[--n_active]; + else active_cursor++; - } } - if (n_in_flight == 0) { - if (n_active == 0) break; + /* ── nothing in flight but buffers remain: relieve CUDA pressure ── */ + if (n_in_flight == 0 && n_active > 0) { io_uring_submit(&fetcher->ioring); -#if defined(__HMLL_CUDA_ENABLED__) - 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++) { - 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; - } - } - } - } -#endif - continue; + if (is_cuda) + hmll_io_uring_cuda_relieve_pressure(fetcher); } - 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 (io_uring_submit_and_wait(&fetcher->ioring, nwait) < 0) { - ctx->error = HMLL_ERR(HMLL_ERR_IO_ERROR); - goto cleanup; + io_uring_submit(&fetcher->ioring); + + /* relieve pressure if submission stalled on busy slots */ + if (n_active > 0 && n_in_flight < HMLL_URING_QUEUE_DEPTH) { + if (is_cuda) hmll_io_uring_cuda_relieve_pressure(fetcher); } - clock_gettime(CLOCK_MONOTONIC, &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) { - for (unsigned i = 0; i < count; i++) { - const struct io_uring_cqe *cqe = cqes[i]; - const uint64_t data = cqe->user_data; - - if (data == HMLL_IO_URING_FADVISE_TAG) continue; - - n_in_flight--; - 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; - - 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]; - void *to = (char *)dsts[bidx].ptr + slot_offsets[slot]; - void *from = fetcher->iovecs[slot].iov_base; - cudaMemcpyAsync(to, from, (size_t)cqe->res, cudaMemcpyHostToDevice, cctx->stream); - cudaEventRecord(cctx->done, cctx->stream); - hmll_io_uring_cuda_stream_set_memcpy(&cctx->state); - } -#else - (void)bidx; -#endif + + /* ── complete: non-blocking peek first, block only when the pipeline full ── */ + unsigned count = io_uring_peek_batch_cqe(&fetcher->ioring, cqes, HMLL_URING_QUEUE_DEPTH); + if (count == 0 && n_in_flight > 0) { + struct io_uring_cqe *cqe; + if (unlikely(io_uring_wait_cqe(&fetcher->ioring, &cqe) < 0)) { + ctx->error = HMLL_ERR(HMLL_ERR_IO_ERROR); + goto cleanup; + } + cqes[0] = cqe; + count = 1; + } + + for (unsigned i = 0; i < count; i++) { + if (cqes[i]->user_data == HMLL_IO_URING_FADVISE_TAG) continue; + + n_in_flight--; + const ssize_t r = hmll_io_uring_fetchv_handle_cqe(ctx, fetcher, cqes[i], dsts, slot_offsets, is_cuda); + if (r < 0) { + io_uring_cq_advance(&fetcher->ioring, count); + goto cleanup; } - io_uring_cq_advance(&fetcher->ioring, count); + nbytes += (size_t)r; } + io_uring_cq_advance(&fetcher->ioring, count); } hmll_io_uring_sync(dsts[0].device, fetcher); @@ -459,14 +610,156 @@ static ssize_t hmll_io_uring_fetchv_impl( return -1; } -struct hmll_error hmll_io_uring_init(struct hmll *ctx, const struct hmll_device device) { - if (hmll_check(ctx->error)) - return ctx->error; +#if defined(__HMLL_CUDA_ENABLED__) +/** + * Decompose buffer @p i into an aligned core and up to 2 edge fragments. + * Issues fadvise for all edge regions on the buffered fd. Returns the number + * of new core/edge entries appended. + */ +static void hmll_io_uring_decompose_buf( + struct hmll_io_uring *fetcher, + const int bfd, + const struct hmll_iobuf *dst, + const size_t off, + struct hmll_iobuf *core_dsts, size_t *core_offs, size_t *n_core, + struct hmll_iobuf *edge_dsts, size_t *edge_offs, size_t *n_edge +) { + const size_t end = off + dst->size; + const size_t a_start = ALIGN_UP(off, ALIGN_PAGE); + const size_t a_end = ALIGN_DOWN(end, ALIGN_PAGE); + const size_t head_sz = a_start - off; + const size_t tail_sz = end - a_end; + const size_t core_sz = (a_end > a_start) ? a_end - a_start : 0; + + if (core_sz < ALIGN_PAGE * 2) { + edge_dsts[*n_edge] = *dst; + edge_offs[*n_edge] = off; + (*n_edge)++; + hmll_io_uring_queue_fadvise(fetcher, bfd, off, dst->size); + return; + } - struct hmll_io_uring *backend = calloc(1, sizeof(struct hmll_io_uring)); - hmll_io_uring_cca_init(&backend->iocca); + core_dsts[*n_core] = (struct hmll_iobuf){ + .size = core_sz, .ptr = (char *)dst->ptr + head_sz, .device = dst->device, + }; + core_offs[*n_core] = a_start; + (*n_core)++; + + if (head_sz > 0) { + hmll_io_uring_queue_fadvise(fetcher, bfd, off, head_sz); + edge_dsts[*n_edge] = (struct hmll_iobuf){ + .size = head_sz, .ptr = dst->ptr, .device = dst->device, + }; + edge_offs[*n_edge] = off; + (*n_edge)++; + } + + if (tail_sz > 0) { + hmll_io_uring_queue_fadvise(fetcher, bfd, a_end, tail_sz); + edge_dsts[*n_edge] = (struct hmll_iobuf){ + .size = tail_sz, .ptr = (char *)dst->ptr + head_sz + core_sz, .device = dst->device, + }; + edge_offs[*n_edge] = a_end; + (*n_edge)++; + } +} + +/** + * CUDA split-I/O fetchv: decomposes each buffer into an aligned core (O_DIRECT) + * and unaligned head/tail edges (buffered). fadvise(WILLNEED) is issued for all + * edge regions before core I/O begins. + */ +static ssize_t hmll_io_uring_fetchv_cuda_split( + struct hmll *ctx, + const int bfd, + const int dfd, + const struct hmll_iobuf *dsts, + const size_t *offsets, + const size_t n +) { + struct hmll_io_uring *fetcher = ctx->fetcher->backend_impl_; + + /* scratch: up to 3 sub-requests per buffer (1 core + 2 edges) */ + const size_t max_parts = n * 3; + const size_t scratch_need = max_parts * (sizeof(struct hmll_iobuf) + sizeof(size_t)); + + _Alignas(16) uint8_t stack_scratch[16384]; + void *scratch_to_free; + uint8_t *scratch = hmll_scratch_alloc(stack_scratch, sizeof(stack_scratch), + scratch_need, &scratch_to_free); + if (!scratch) { + ctx->error = HMLL_ERR(HMLL_ERR_ALLOCATION_FAILED); + return -1; + } + + struct hmll_iobuf *core_dsts = (struct hmll_iobuf *)scratch; + size_t *core_offs = (size_t *)(scratch + max_parts * sizeof(struct hmll_iobuf)); + struct hmll_iobuf *edge_dsts = core_dsts + n; + size_t *edge_offs = core_offs + n; + size_t n_core = 0, n_edge = 0; + + for (size_t i = 0; i < n; i++) { + if (dsts[i].size == 0) continue; + hmll_io_uring_decompose_buf(fetcher, bfd, &dsts[i], offsets[i], + core_dsts, core_offs, &n_core, + edge_dsts, edge_offs, &n_edge); + } + + io_uring_submit(&fetcher->ioring); + + ssize_t total = 0; + + if (n_core > 0) { + ssize_t r = hmll_io_uring_fetchv_loop(ctx, dfd, core_dsts, core_offs, n_core, 0); + if (r < 0) goto cleanup; + total += r; + } + + if (n_edge > 0) { + ssize_t r = hmll_io_uring_fetchv_loop(ctx, bfd, edge_dsts, edge_offs, n_edge, 0); + if (r < 0) goto cleanup; + total += r; + } + + if (scratch_to_free) free(scratch_to_free); + return total; + +cleanup: + if (scratch_to_free) free(scratch_to_free); + return -1; +} +#endif /* __HMLL_CUDA_ENABLED__ */ + +static ssize_t hmll_io_uring_fetchv_impl( + struct hmll *ctx, + const unsigned iofile, + const struct hmll_iobuf *dsts, + const size_t *offsets, + const size_t n +) { + if (hmll_check(ctx->error)) return -1; + if (unlikely(n == 0)) return 0; + + const unsigned bfd = hmll_io_uring_buffered_fd(iofile); + +#if defined(__HMLL_CUDA_ENABLED__) + if (hmll_device_is_cuda(dsts[0].device)) { + const int dfd = hmll_io_uring_direct_fd(iofile); + return hmll_io_uring_fetchv_cuda_split(ctx, bfd, dfd, dsts, offsets, n); + } +#endif + return hmll_io_uring_fetchv_loop(ctx, bfd, dsts, offsets, n, 1); +} + +static struct hmll_error hmll_io_uring_queue_init( + struct hmll *ctx, + struct hmll_io_uring *backend, + const struct hmll_device device +) { + (void)ctx; struct io_uring_params params = { + .sq_thread_cpu = 0, .flags = hmll_io_uring_get_setup_flags(), .sq_thread_idle = 500 }; @@ -475,8 +768,7 @@ struct hmll_error hmll_io_uring_init(struct hmll *ctx, const struct hmll_device #if defined(__HMLL_CUDA_ENABLED__) cudaError_t cuda_err = cudaSetDevice(device.idx); if (cuda_err != cudaSuccess) { - ctx->error = HMLL_ERR(HMLL_ERR_CUDA_SET_DEVICE_FAILED); - return ctx->error; + return HMLL_ERR(HMLL_ERR_CUDA_SET_DEVICE_FAILED); } struct hmll_io_uring_cuda_context *data = calloc(HMLL_URING_QUEUE_DEPTH, sizeof(struct hmll_io_uring_cuda_context)); @@ -488,34 +780,61 @@ struct hmll_error hmll_io_uring_init(struct hmll *ctx, const struct hmll_device CHECK_CUDA(cudaEventCreateWithFlags(&data[i].done, cudaEventDisableTiming)); } - int res = 0; - if ((res = io_uring_queue_init_params(HMLL_URING_QUEUE_DEPTH, &backend->ioring, ¶ms)) < 0) { - ctx->error = HMLL_SYS_ERR(-res); - return ctx->error; + // we get the "optimal" set of flags we would like to enable and attempt to init there + int res = io_uring_queue_init_params(HMLL_URING_QUEUE_DEPTH, &backend->ioring, ¶ms); + if (res < 0) { + if (res == -EINVAL && (params.flags & IORING_SETUP_COOP_TASKRUN)) { + params.flags = HMLL_IO_URING_DEFAULT_QUEUE_PARAMS; + res = io_uring_queue_init_params(HMLL_URING_QUEUE_DEPTH, &backend->ioring, ¶ms); + } + if (res < 0) { + return HMLL_SYS_ERR(-res); + } } - ctx->error = hmll_io_uring_register_staging_buffers(ctx, backend, device); - if (hmll_check(ctx->error)) { - return ctx->error; + struct hmll_error err = hmll_io_uring_register_staging_buffers(ctx, backend, device); + if (hmll_check(err)) { + return err; } + return HMLL_OK; #else - ctx->error = HMLL_ERR(HMLL_ERR_CUDA_NOT_ENABLED); - return ctx->error; + return HMLL_ERR(HMLL_ERR_CUDA_NOT_ENABLED); #endif } else { - int res; - if ((res = io_uring_queue_init_params(HMLL_URING_QUEUE_DEPTH, &backend->ioring, ¶ms)) < 0) { - ctx->error = HMLL_SYS_ERR(-res); - goto cleanup; + int res = io_uring_queue_init_params(HMLL_URING_QUEUE_DEPTH, &backend->ioring, ¶ms); + if (res < 0) { + if (res == -EINVAL && (params.flags & IORING_SETUP_COOP_TASKRUN)) { + params.flags = IORING_SETUP_SQPOLL; + res = io_uring_queue_init_params(HMLL_URING_QUEUE_DEPTH, &backend->ioring, ¶ms); + } + if (res < 0) { + return HMLL_SYS_ERR(-res); + } } + return HMLL_OK; } +} + +struct hmll_error hmll_io_uring_init(struct hmll *ctx, const struct hmll_device device) { + if (hmll_check(ctx->error)) + return ctx->error; + + struct hmll_io_uring *backend = calloc(1, sizeof(struct hmll_io_uring)); - int *iofiles = calloc(ctx->num_sources, sizeof(int)); - for (size_t i = 0; i < ctx->num_sources; ++i) - iofiles[i] = ctx->sources[i].fd; + ctx->error = hmll_io_uring_queue_init(ctx, backend, device); + if (hmll_check(ctx->error)) + goto cleanup; + + const size_t n_iofiles = ctx->num_sources * 2; + int *iofiles = calloc(n_iofiles, sizeof(int)); + for (unsigned i = 0; i < ctx->num_sources; ++i) { + iofiles[hmll_io_uring_buffered_fd(i)] = ctx->sources[i].fd; + const int dfd = ctx->sources[i].d_fd; + iofiles[hmll_io_uring_direct_fd(i)] = (dfd > 0) ? dfd : ctx->sources[i].fd; + } - const int res = io_uring_register_files(&backend->ioring, iofiles, ctx->num_sources); + const int res = io_uring_register_files(&backend->ioring, iofiles, n_iofiles); free(iofiles); if (res != 0) { diff --git a/lib/python/loader.cpp b/lib/python/loader.cpp index baab8d5..de38caf 100644 --- a/lib/python/loader.cpp +++ b/lib/python/loader.cpp @@ -1,12 +1,15 @@ #include "loader.hpp" #include +#include +#include #include #include + +#include "hmll/memory.h" #include "formatters.hpp" + #include "ndarray.hpp" -#include "fmt/compile.h" -#include "hmll/memory.h" namespace nb = nanobind; using namespace nb::literals; @@ -146,52 +149,6 @@ size_t WeightLoader::fetchv(const int iofile, const std::vector(m, "Device", R"pbdoc(Define all the targetable devices)pbdoc") - .def_static("cpu", &hmll_device_cpu, "Create CPU device") - .def_static("cuda", &hmll_device_cuda, "idx"_a = 0, "Create CUDA device with index") - .def_prop_ro("kind", [](const hmll_device_t& d) { return d.kind; }) - .def_prop_ro("idx", [](const hmll_device_t& d) { return d.idx; }) - .def_prop_ro("is_cpu", [](const hmll_device_t& d) { return hmll_device_is_cpu(d); }) - .def_prop_ro("is_cuda", [](const hmll_device_t& d) { return hmll_device_is_cuda(d); }) - .def("__eq__", [](const hmll_device_t& a, const hmll_device_t& b) { return hmll_device_eq(a, b); }) - .def("__repr__", [](const hmll_device_t& d) { return hmll_device_is_cpu(d) ? "Device.cpu()" : fmt::format("Device.cuda({})", d.idx); }); - - nb::enum_(m, "DeviceKind", R"pbdoc(Define all the targetable devices)pbdoc") - .value("CPU", HMLL_DEVICE_CPU) - .value("CUDA", HMLL_DEVICE_CUDA); - - nb::enum_(m, "Backend", R"pbdoc(Define the I/O backend to use)pbdoc") - .value("AUTO", HMLL_FETCHER_AUTO, "Automatically select backend (defaults to MMAP)") -#ifdef __HMLL_IO_URING_ENABLED__ - .value("IO_URING", HMLL_FETCHER_IO_URING, "Use io_uring for async I/O (Linux only)") -#endif - .value("MMAP", HMLL_FETCHER_MMAP, "Use memory-mapped I/O"); - - nb::enum_(m, "dtype", R"pbdoc(Define all the targetable element type in a tensor)pbdoc") - .value("BOOL", HMLL_DTYPE_BOOL) - .value("BFLOAT16", HMLL_DTYPE_BFLOAT16) - .value("COMPLEX", HMLL_DTYPE_COMPLEX) - .value("FLOAT4", HMLL_DTYPE_FLOAT4) - .value("FLOAT6_E2M3", HMLL_DTYPE_FLOAT6_E2M3) - .value("FLOAT6_E3M2", HMLL_DTYPE_FLOAT6_E3M2) - .value("FLOAT8_E5M2", HMLL_DTYPE_FLOAT8_E5M2) - .value("FLOAT8_E4M3", HMLL_DTYPE_FLOAT8_E4M3) - .value("FLOAT8_E8M0", HMLL_DTYPE_FLOAT8_E8M0) - .value("FLOAT16", HMLL_DTYPE_FLOAT16) - .value("FLOAT32", HMLL_DTYPE_FLOAT32) - .value("FLOAT64", HMLL_DTYPE_FLOAT64) - .value("SIGNED_INT4", HMLL_DTYPE_SIGNED_INT4) - .value("SIGNED_INT8", HMLL_DTYPE_SIGNED_INT8) - .value("SIGNED_INT16", HMLL_DTYPE_SIGNED_INT16) - .value("SIGNED_INT32", HMLL_DTYPE_SIGNED_INT32) - .value("SIGNED_INT64", HMLL_DTYPE_SIGNED_INT64) - .value("UNSIGNED_INT4", HMLL_DTYPE_UNSIGNED_INT4) - .value("UNSIGNED_INT8", HMLL_DTYPE_UNSIGNED_INT8) - .value("UNSIGNED_INT16", HMLL_DTYPE_UNSIGNED_INT16) - .value("UNSIGNED_INT32", HMLL_DTYPE_UNSIGNED_INT32) - .value("UNSIGNED_INT64", HMLL_DTYPE_UNSIGNED_INT64) - .value("UNKNOWN", HMLL_DTYPE_UNKNOWN); - nb::class_(m, "WeightLoader", R"pbdoc("Opaque type representing an allocated fetcher backend)pbdoc") .def(nb::new_(&WeightLoader::from_paths), "paths"_a.sig("list[str]"), "device"_a.sig("Device")) .def_prop_ro("device", &WeightLoader::device) @@ -201,6 +158,6 @@ void init_loader(nb::module_& m) .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()); + return fmt::format(FMT_COMPILE("WeightLoader(kind={}, device={})"), self.kind(), self.device()); }); } diff --git a/lib/python/pyhmll.cpp b/lib/python/pyhmll.cpp index ad963ef..fbe0586 100644 --- a/lib/python/pyhmll.cpp +++ b/lib/python/pyhmll.cpp @@ -1,7 +1,10 @@ -#include +#include #include +#include "hmll/hmll.h" namespace nb = nanobind; +using namespace nb::literals; + void init_loader(nb::module_&); @@ -14,6 +17,52 @@ NB_MODULE(_pyhmll_impl, m) { m.doc() = "hmll: High-Performance Model Loading Library - Efficient AI Model loading for modern AI workloads."; + nb::class_(m, "Device", R"pbdoc(Define all the targetable devices)pbdoc") + .def_static("cpu", &hmll_device_cpu, "Create CPU device") + .def_static("cuda", &hmll_device_cuda, "idx"_a.sig("int = 0"), "Create CUDA device with index") + .def_prop_ro("kind", [](const hmll_device_t& d) { return d.kind; }) + .def_prop_ro("idx", [](const hmll_device_t& d) { return d.idx; }) + .def_prop_ro("is_cpu", [](const hmll_device_t& d) { return hmll_device_is_cpu(d); }) + .def_prop_ro("is_cuda", [](const hmll_device_t& d) { return hmll_device_is_cuda(d); }) + .def("__eq__", [](const hmll_device_t& a, const hmll_device_t& b) { return hmll_device_eq(a, b); }) + .def("__repr__", [](const hmll_device_t& d) { return hmll_device_is_cpu(d) ? "Device.cpu()" : fmt::format("Device.cuda({})", d.idx); }); + + nb::enum_(m, "DeviceKind", R"pbdoc(Define all the targetable devices)pbdoc") + .value("CPU", HMLL_DEVICE_CPU) + .value("CUDA", HMLL_DEVICE_CUDA); + + nb::enum_(m, "Backend", R"pbdoc(Define the I/O backend to use)pbdoc") + .value("AUTO", HMLL_FETCHER_AUTO, "Automatically select backend (defaults to MMAP)") +#ifdef __HMLL_IO_URING_ENABLED__ + .value("IO_URING", HMLL_FETCHER_IO_URING, "Use io_uring for async I/O (Linux only)") +#endif + .value("MMAP", HMLL_FETCHER_MMAP, "Use memory-mapped I/O"); + + nb::enum_(m, "dtype", R"pbdoc(Define all the targetable element type in a tensor)pbdoc") + .value("BOOL", HMLL_DTYPE_BOOL) + .value("BFLOAT16", HMLL_DTYPE_BFLOAT16) + .value("COMPLEX", HMLL_DTYPE_COMPLEX) + .value("FLOAT4", HMLL_DTYPE_FLOAT4) + .value("FLOAT6_E2M3", HMLL_DTYPE_FLOAT6_E2M3) + .value("FLOAT6_E3M2", HMLL_DTYPE_FLOAT6_E3M2) + .value("FLOAT8_E5M2", HMLL_DTYPE_FLOAT8_E5M2) + .value("FLOAT8_E4M3", HMLL_DTYPE_FLOAT8_E4M3) + .value("FLOAT8_E8M0", HMLL_DTYPE_FLOAT8_E8M0) + .value("FLOAT16", HMLL_DTYPE_FLOAT16) + .value("FLOAT32", HMLL_DTYPE_FLOAT32) + .value("FLOAT64", HMLL_DTYPE_FLOAT64) + .value("SIGNED_INT4", HMLL_DTYPE_SIGNED_INT4) + .value("SIGNED_INT8", HMLL_DTYPE_SIGNED_INT8) + .value("SIGNED_INT16", HMLL_DTYPE_SIGNED_INT16) + .value("SIGNED_INT32", HMLL_DTYPE_SIGNED_INT32) + .value("SIGNED_INT64", HMLL_DTYPE_SIGNED_INT64) + .value("UNSIGNED_INT4", HMLL_DTYPE_UNSIGNED_INT4) + .value("UNSIGNED_INT8", HMLL_DTYPE_UNSIGNED_INT8) + .value("UNSIGNED_INT16", HMLL_DTYPE_UNSIGNED_INT16) + .value("UNSIGNED_INT32", HMLL_DTYPE_UNSIGNED_INT32) + .value("UNSIGNED_INT64", HMLL_DTYPE_UNSIGNED_INT64) + .value("UNKNOWN", HMLL_DTYPE_UNKNOWN); + init_loader(m); #ifdef __HMLL_SAFETENSORS_ENABLED__ diff --git a/lib/python/pyhmll/torch.py b/lib/python/pyhmll/torch.py index 43271ea..3d15f3f 100644 --- a/lib/python/pyhmll/torch.py +++ b/lib/python/pyhmll/torch.py @@ -71,8 +71,8 @@ def device_to_hmll(device: torch.device) -> Device: """ match device.type: case "cuda": - return Device.CUDA + return Device.cuda(device.index) case "cpu": - return Device.CPU + return Device.cpu() case _: raise ValueError(f"Unsupported device for pyhmll: {device!r}") \ No newline at end of file diff --git a/lib/rust/hmll-sys/src/lib.rs b/lib/rust/hmll-sys/src/lib.rs index 3c6ec68..a6487ce 100644 --- a/lib/rust/hmll-sys/src/lib.rs +++ b/lib/rust/hmll-sys/src/lib.rs @@ -170,11 +170,13 @@ mod tests { fn test_source_size() { let source = hmll_source { fd: -1, + d_fd: -1, size: 1024, content: std::ptr::null(), }; assert_eq!(source.size, 1024); assert_eq!(source.fd, -1); + assert_eq!(source.d_fd, -1); assert!(source.content.is_null()); } diff --git a/lib/rust/hmll/src/source.rs b/lib/rust/hmll/src/source.rs index b826cf9..1d9d6fa 100644 --- a/lib/rust/hmll/src/source.rs +++ b/lib/rust/hmll/src/source.rs @@ -60,6 +60,7 @@ impl Source { let mut source = hmll_sys::hmll_source { fd: -1, + d_fd: -1, size: 0, content: null_mut(), }; diff --git a/lib/safetensors.c b/lib/safetensors.c index cb4c3fa..a7a991c 100644 --- a/lib/safetensors.c +++ b/lib/safetensors.c @@ -177,8 +177,8 @@ size_t hmll_safetensors_index(struct hmll *ctx, struct hmll_registry *reg, const goto cleanup; } - yyjson_val *root = yyjson_doc_get_root(document); - yyjson_val *map = yyjson_obj_get(root, "weight_map"); + const yyjson_val *root = yyjson_doc_get_root(document); + const yyjson_val *map = yyjson_obj_get(root, "weight_map"); if (map == NULL) { ctx->error = HMLL_ERR(HMLL_ERR_SAFETENSORS_JSON_MALFORMED_INDEX); goto cleanup; @@ -256,13 +256,13 @@ size_t hmll_safetensors_populate_registry( ) { size_t tidx = 0, num_tensors = 0; if (hmll_check(ctx->error)) - goto exit; + return 0; yyjson_doc *document = NULL; FILE *file = hmll_get_file_from_fd(source); if (!file) { ctx->error = HMLL_ERR(HMLL_ERR_FILE_OPEN_FAILED); - goto freeup_and_exit; + goto freeup; } uint64_t hsize; @@ -273,13 +273,13 @@ size_t hmll_safetensors_populate_registry( document = yyjson_read_opts((char *)source.content + sizeof(uint64_t), hsize, YYJSON_READ_NOFLAG, NULL, &error); if (!document) { ctx->error = HMLL_ERR(HMLL_ERR_SAFETENSORS_JSON_INVALID_HEADER); - goto freeup_and_exit; + goto freeup; } - yyjson_val *root = yyjson_doc_get_root(document); + const yyjson_val *root = yyjson_doc_get_root(document); if (!yyjson_is_obj(root)) { ctx->error = HMLL_ERR(HMLL_ERR_SAFETENSORS_JSON_INVALID_HEADER); - goto freeup_and_exit; + goto freeup; } // we don't allocate this the number of tensors is already set on the context @@ -289,17 +289,17 @@ size_t hmll_safetensors_populate_registry( if (reg->num_tensors == 0) { if ((reg->names = calloc(num_tensors, sizeof(char*))) == NULL) { ctx->error = HMLL_ERR(HMLL_ERR_ALLOCATION_FAILED); - goto freeup_and_exit; + goto freeup; } if ((reg->tensors = calloc(num_tensors, sizeof(struct hmll_tensor_specs))) == NULL) { ctx->error = HMLL_ERR(HMLL_ERR_ALLOCATION_FAILED); - goto freeup_and_exit; + goto freeup; } if ((reg->indexes = calloc(num_tensors, sizeof(struct hmll_source))) == NULL) { ctx->error = HMLL_ERR(HMLL_ERR_ALLOCATION_FAILED); - goto freeup_and_exit; + goto freeup; } } @@ -323,12 +323,12 @@ size_t hmll_safetensors_populate_registry( if (!yyjson_is_obj(val)) { ctx->error = HMLL_ERR(HMLL_ERR_SAFETENSORS_JSON_MALFORMED_HEADER); - goto freeup_and_exit; + goto freeup; } if (hmll_check(hmll_safetensors_header_parse_tensor(val, tensors + offset + tidx))) { ctx->error = HMLL_ERR(HMLL_ERR_SAFETENSORS_JSON_MALFORMED_HEADER); - goto freeup_and_exit; + goto freeup;; } // tensor offsets start at 0, we need to add header size + 8 to get the real position in the file @@ -338,11 +338,24 @@ size_t hmll_safetensors_populate_registry( ++tidx; } -freeup_and_exit: + goto cleanup; + +freeup: + if (reg->names) { + for (size_t i = 0; i < tidx; ++i) { + free(reg->names[i]); + } + free(reg->names); + reg->names = NULL; + } + free(reg->tensors); + reg->tensors = NULL; + free(reg->indexes); + reg->indexes = NULL; + +cleanup: if (document) yyjson_doc_free(document); -exit: - if (hmll_check(ctx->error)) return 0; if (reg->num_tensors == 0) reg->num_tensors = tidx; return tidx; } diff --git a/lib/tensors.c b/lib/tensors.c index 5c96cc2..14aa21d 100644 --- a/lib/tensors.c +++ b/lib/tensors.c @@ -118,4 +118,4 @@ size_t hmll_numel(const struct hmll_tensor_specs *specs) numel *= specs->shape[i]; return numel; -} \ No newline at end of file +} diff --git a/lib/unix/backend/mmap.c b/lib/unix/backend/mmap.c index 109b1ba..c9ecc24 100644 --- a/lib/unix/backend/mmap.c +++ b/lib/unix/backend/mmap.c @@ -21,7 +21,7 @@ #endif static ssize_t -hmll_mmap_fetch_range_impl(struct hmll *ctx, const int iofile, const struct hmll_iobuf *dst, const size_t offset) +hmll_mmap_fetch_range_impl(struct hmll *ctx, const unsigned iofile, const struct hmll_iobuf *dst, const size_t offset) { if (hmll_check(ctx->error)) return -1; if (dst->size == 0) return 0; @@ -44,7 +44,7 @@ hmll_mmap_fetch_range_impl(struct hmll *ctx, const int iofile, const struct hmll } static ssize_t -hmll_mmap_fetchv_range_impl(struct hmll *ctx, const int iofile, const struct hmll_iobuf *dsts, const size_t *offsets, const size_t n) +hmll_mmap_fetchv_range_impl(struct hmll *ctx, const unsigned iofile, const struct hmll_iobuf *dsts, const size_t *offsets, const size_t n) { if (hmll_check(ctx->error)) return -1; diff --git a/lib/unix/file.c b/lib/unix/file.c index c8566e4..9140adb 100644 --- a/lib/unix/file.c +++ b/lib/unix/file.c @@ -43,6 +43,12 @@ struct hmll_error hmll_source_open(const char *path, struct hmll_source *src) src->size = sb.st_size; src->content = content; +#if defined(__linux__) + src->d_fd = open(path, O_RDONLY | O_DIRECT); +#else + src->d_fd = -1; +#endif + return HMLL_OK; close_fd_then_exit: @@ -55,9 +61,13 @@ struct hmll_error hmll_source_open(const char *path, struct hmll_source *src) void hmll_source_close(struct hmll_source *src) { - if (src && src->fd > 0) { + if (src && src->fd != -1) { close(src->fd); - src->fd = -1; // Mark as closed + src->fd = -1; + } + if (src && src->d_fd != -1) { + close(src->d_fd); + src->d_fd = -1; } } diff --git a/scripts/create_fetchv_testing_safetensors.py b/scripts/create_fetchv_testing_safetensors.py new file mode 100644 index 0000000..d524d9e --- /dev/null +++ b/scripts/create_fetchv_testing_safetensors.py @@ -0,0 +1,172 @@ +""" +Generate safetensors fixtures for hmll_fetchv tests. + +Produces: +1. fetchv_test.safetensors - the single file (tp=1) with deterministic fill patterns, + multiple dtypes, various sizes (including >512KB for io_uring chunking). +2. fetchv_sharded/ - directory with model.safetensors.index.json and shard files + (model-00001-of-00003.safetensors, etc.) for distributed/TP-style tests. + +Fill pattern: value[i] = f(i) so C++ tests can validate byte-level correctness. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import torch +from safetensors.torch import save_file + +# HMLL_URING_BUFFER_SIZE is 512*1024 bytes; use tensor larger than that for chunking tests +IO_URING_BUFFER_BYTES = 512 * 1024 +LARGE_F32_NUMEL = (IO_URING_BUFFER_BYTES // 4) + 1024 # >512KB in float32 + + +def make_deterministic_tensor(dtype: torch.dtype, shape: list[int], name: str) -> torch.Tensor: + """Create a tensor with deterministic values: value[i] = f(i) for validation.""" + numel = 1 + for s in shape: + numel *= s + if numel == 0: + numel = 1 + + if dtype in (torch.float32, torch.float16, torch.bfloat16): + t = torch.arange(numel, dtype=torch.float32) + if dtype != torch.float32: + t = t.to(dtype) + return t.reshape(shape).contiguous() + if dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + # float8 has limited range; use small integers + t = torch.arange(min(numel, 256), dtype=torch.float32) + t = t.to(dtype) + if numel > 256: + t = t.repeat((numel + 255) // 256)[:numel] + return t.reshape(shape).contiguous() + if dtype in (torch.int8, torch.int16, torch.int32, torch.int64): + t = torch.arange(numel, dtype=torch.int64) + t = (t % (1 << 15)) - (1 << 14) # spread in signed range + return t.to(dtype).reshape(shape).contiguous() + if dtype in (torch.uint8, torch.uint16, torch.uint32, torch.uint64): + t = torch.arange(numel, dtype=torch.int64) + t = t % (1 << 8) if dtype == torch.uint8 else t % (1 << 16) + return t.to(dtype).reshape(shape).contiguous() + if dtype == torch.complex64: + r = torch.arange(numel, dtype=torch.float32) + i = torch.arange(numel, dtype=torch.float32) * 0.5 + return (torch.complex(r, i)).reshape(shape).contiguous() + raise ValueError(f"Unsupported dtype for fetchv tests: {dtype}") + + +def build_single_file_tensors() -> dict[str, torch.Tensor]: + """Tensors for single-file fixture: multiple dtypes, sizes, deterministic data.""" + tensors: dict[str, torch.Tensor] = {} + + # Scalars (rank-0) + for dtype in [torch.float32, torch.int32, torch.uint8]: + name = str(dtype).replace("torch.", "") + ".scalar" + tensors[name] = make_deterministic_tensor(dtype, [], name) + + # 1D vectors: small (16), medium (1024), large (8192) + for dtype in [torch.float32, torch.float16, torch.bfloat16, torch.int32, torch.int64, torch.uint8]: + base = str(dtype).replace("torch.", "") + for size, suffix in [(16, "vec16"), (1024, "vec1024"), (8192, "vec8192")]: + tensors[f"{base}.{suffix}"] = make_deterministic_tensor(dtype, [size], f"{base}.{suffix}") + + # Large tensor > 512KB to exercise io_uring chunked path + tensors["float32.large"] = make_deterministic_tensor(torch.float32, [LARGE_F32_NUMEL], "float32.large") + tensors["int32.large"] = make_deterministic_tensor( + torch.int32, [LARGE_F32_NUMEL], "int32.large" + ) + + # 2D matrices + for dtype in [torch.float32, torch.int32]: + base = str(dtype).replace("torch.", "") + tensors[f"{base}.mat64"] = make_deterministic_tensor(dtype, [64, 64], f"{base}.mat64") + + # float8 for mixed-dtype fetchv tests + tensors["float8_e4m3fn.vec16"] = make_deterministic_tensor( + torch.float8_e4m3fn, [16], "float8_e4m3fn.vec16" + ) + tensors["float8_e5m2.vec16"] = make_deterministic_tensor( + torch.float8_e5m2, [16], "float8_e5m2.vec16" + ) + + return tensors + + +def build_sharded_tensors() -> tuple[dict[str, dict[str, torch.Tensor]], dict[str, str]]: + """ + Build tensors split across 3 shards and weight_map for index. + Returns (shard_name -> {tensor_name -> tensor}, weight_map). + """ + num_shards = 3 + shard_names = [f"model-{i:05}-of-{num_shards:05}.safetensors" for i in range(1, num_shards + 1)] + shards: dict[str, dict[str, torch.Tensor]] = {s: {} for s in shard_names} + weight_map: dict[str, str] = {} + + # Shard 0: float32 and float16 + for size, suffix in [(16, "vec16"), (256, "vec256")]: + shards[shard_names[0]][f"float32.shard0.{suffix}"] = make_deterministic_tensor( + torch.float32, [size], f"float32.shard0.{suffix}" + ) + weight_map[f"float32.shard0.{suffix}"] = shard_names[0] + shards[shard_names[0]][f"float16.shard0.{suffix}"] = make_deterministic_tensor( + torch.float16, [size], f"float16.shard0.{suffix}" + ) + weight_map[f"float16.shard0.{suffix}"] = shard_names[0] + + # Shard 1: int32, int64, uint8 + for size, suffix in [(16, "vec16"), (128, "vec128")]: + shards[shard_names[1]][f"int32.shard1.{suffix}"] = make_deterministic_tensor( + torch.int32, [size], f"int32.shard1.{suffix}" + ) + weight_map[f"int32.shard1.{suffix}"] = shard_names[1] + shards[shard_names[1]][f"uint8.shard1.{suffix}"] = make_deterministic_tensor( + torch.uint8, [size], f"uint8.shard1.{suffix}" + ) + weight_map[f"uint8.shard1.{suffix}"] = shard_names[1] + shards[shard_names[1]]["int64.shard1.scalar"] = make_deterministic_tensor( + torch.int64, [], "int64.shard1.scalar" + ) + weight_map["int64.shard1.scalar"] = shard_names[1] + + # Shard 2: bfloat16 and scalars + shards[shard_names[2]]["bfloat16.shard2.vec64"] = make_deterministic_tensor( + torch.bfloat16, [64], "bfloat16.shard2.vec64" + ) + weight_map["bfloat16.shard2.vec64"] = shard_names[2] + shards[shard_names[2]]["float32.shard2.scalar"] = make_deterministic_tensor( + torch.float32, [], "float32.shard2.scalar" + ) + weight_map["float32.shard2.scalar"] = shard_names[2] + + return shards, weight_map + + +def main() -> None: + cwd = Path(os.getcwd()) + + # 1) Single-file fixture + single_tensors = build_single_file_tensors() + single_path = cwd / "fetchv_test.safetensors" + save_file(single_tensors, single_path) + print(single_path) + + # 2) Sharded fixtures + shard_dir = cwd / "fetchv_sharded" + shard_dir.mkdir(exist_ok=True) + shards, weight_map = build_sharded_tensors() + for shard_name, tensors in shards.items(): + if tensors: + save_file(tensors, shard_dir / shard_name) + index = {"metadata": {"format": "pt"}, "weight_map": weight_map} + index_path = shard_dir / "model.safetensors.index.json" + with open(index_path, "w") as f: + json.dump(index, f, indent=2) + print(index_path) + + +if __name__ == "__main__": + main() diff --git a/tests/tests_hmll_fetchv.cpp b/tests/tests_hmll_fetchv.cpp index 2af8607..d70df75 100644 --- a/tests/tests_hmll_fetchv.cpp +++ b/tests/tests_hmll_fetchv.cpp @@ -77,7 +77,7 @@ TEST_CASE("fetchv - single-element fetchv (n=1)", "[fetchv][safetensors]") { for (const auto& [name, backend] : kBackends) { INFO("Backend: " << name); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + 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); @@ -85,7 +85,7 @@ TEST_CASE("fetchv - single-element fetchv (n=1)", "[fetchv][safetensors]") { 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); + hmll_iobuf_t buffer = hmll_get_buffer_for_range(&ctx, range); REQUIRE_FALSE(hmll_check(ctx.error)); size_t offsets[1] = {range.start}; @@ -114,7 +114,7 @@ TEST_CASE("fetchv - multi-element same dtype", "[fetchv][safetensors]") { for (const auto& [name, backend] : kBackends) { INFO("Backend: " << name); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + 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]; @@ -124,7 +124,7 @@ TEST_CASE("fetchv - multi-element same dtype", "[fetchv][safetensors]") { 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); + dsts[i] = hmll_get_buffer_for_range(&ctx, range); REQUIRE_FALSE(hmll_check(ctx.error)); offsets[i] = range.start; total += dsts[i].size; @@ -157,7 +157,7 @@ TEST_CASE("fetchv - multi-element mixed dtypes", "[fetchv][safetensors]") { for (const auto& [name, backend] : kBackends) { INFO("Backend: " << name); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + 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"); @@ -165,9 +165,9 @@ TEST_CASE("fetchv - multi-element mixed dtypes", "[fetchv][safetensors]") { 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}), + hmll_get_buffer_for_range(&ctx, {l_f32.specs->start, l_f32.specs->end}), + hmll_get_buffer_for_range(&ctx, {l_i32.specs->start, l_i32.specs->end}), + hmll_get_buffer_for_range(&ctx, {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)); @@ -198,7 +198,7 @@ TEST_CASE("fetchv - scattered reads within single tensor", "[fetchv][safetensors for (const auto& [name, backend] : kBackends) { INFO("Backend: " << name); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + 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); @@ -216,7 +216,7 @@ TEST_CASE("fetchv - scattered reads within single tensor", "[fetchv][safetensors 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); + dsts[i] = hmll_get_buffer_for_range(&ctx, r); offsets[i] = r.start; } @@ -245,13 +245,13 @@ TEST_CASE("fetchv - full tensor via fetchv matches hmll_fetch", "[fetchv][safete for (const auto& [name, backend] : kBackends) { INFO("Backend: " << name); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + 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); + hmll_iobuf_t buf_fetch = hmll_get_buffer_for_range(&ctx, range); + hmll_iobuf_t buf_fetchv = hmll_get_buffer_for_range(&ctx, range); REQUIRE_FALSE(hmll_check(ctx.error)); ssize_t r1 = hmll_fetch(&ctx, lookup.file, &buf_fetch, range.start); @@ -276,20 +276,24 @@ 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))); + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); - ssize_t ret = hmll_fetchv(&ctx, 0, nullptr, nullptr, 0); - REQUIRE(ret == 0); - REQUIRE_FALSE(hmll_check(ctx.error)); + 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(), backend))); - hmll_free_registry(®istry); - hmll_destroy(&ctx); - hmll_source_close(&src); + 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]") { @@ -304,7 +308,7 @@ TEST_CASE("fetchv - scalar tensors", "[fetchv][safetensors]") { for (const auto& [name, backend] : kBackends) { INFO("Backend: " << name); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + 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"); @@ -312,9 +316,9 @@ TEST_CASE("fetchv - scalar tensors", "[fetchv][safetensors]") { 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}), + hmll_get_buffer_for_range(&ctx, {l0.specs->start, l0.specs->end}), + hmll_get_buffer_for_range(&ctx, {l1.specs->start, l1.specs->end}), + hmll_get_buffer_for_range(&ctx, {l2.specs->start, l2.specs->end}), }; size_t offsets[3] = {l0.specs->start, l1.specs->start, l2.specs->start}; @@ -344,12 +348,12 @@ TEST_CASE("fetchv - large tensor exceeds io_uring buffer", "[fetchv][safetensors for (const auto& [name, backend] : kBackends) { INFO("Backend: " << name); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + 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); + hmll_iobuf_t buffer = hmll_get_buffer_for_range(&ctx, range); REQUIRE_FALSE(hmll_check(ctx.error)); REQUIRE(buffer.size > 512 * 1024u); // > HMLL_URING_BUFFER_SIZE @@ -380,7 +384,7 @@ TEST_CASE("fetchv - many concurrent ranges", "[fetchv][safetensors]") { for (const auto& [name, backend] : kBackends) { INFO("Backend: " << name); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, hmll_device_cpu(), backend))); const size_t N = 40; std::vector lookups(N); @@ -390,7 +394,7 @@ TEST_CASE("fetchv - many concurrent ranges", "[fetchv][safetensors]") { 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); + dsts[i] = hmll_get_buffer_for_range(&ctx, range); offsets[i] = range.start; } @@ -411,61 +415,69 @@ 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 (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); - for (auto & dst : dsts) hmll_free_buffer(&dst); - hmll_free_registry(®istry); - hmll_destroy(&ctx); - hmll_source_close(&src); + 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(), backend))); + + 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, 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); + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + + 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(), backend))); + + 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, range); + hmll_iobuf_t buf2 = hmll_get_buffer_for_range(&ctx, 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 --- @@ -519,11 +531,11 @@ TEST_CASE("fetchv - sharded fetchv across files", "[fetchv][safetensors][sharded 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) { + INFO("Backend: " << name); hmll_destroy(&ctx); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, sources.data(), num_files, HMLL_DEVICE_CPU, backend))); + 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"); @@ -536,9 +548,9 @@ TEST_CASE("fetchv - sharded fetchv across files", "[fetchv][safetensors][sharded 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}), + hmll_get_buffer_for_range(&ctx, {l0.specs->start, l0.specs->end}), + hmll_get_buffer_for_range(&ctx, {l1.specs->start, l1.specs->end}), + hmll_get_buffer_for_range(&ctx, {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}; @@ -574,7 +586,7 @@ TEST_CASE("fetchv - fetchv matches fetch for several tensors", "[fetchv][safeten for (const auto& [name, backend] : kBackends) { INFO("Backend: " << name); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + 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; @@ -586,8 +598,8 @@ TEST_CASE("fetchv - fetchv matches fetch for several tensors", "[fetchv][safeten 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); + buf_fetch[i] = hmll_get_buffer_for_range(&ctx, r); + buf_fetchv[i] = hmll_get_buffer_for_range(&ctx, r); offsets[i] = r.start; total += buf_fetch[i].size; } @@ -615,25 +627,29 @@ 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))); + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); - // Poison the context with an error - ctx.error = HMLL_ERR(HMLL_ERR_IO_ERROR); + 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(), backend))); - hmll_iobuf_t dsts[1] = {}; - size_t offsets[1] = {0}; - ssize_t ret = hmll_fetchv(&ctx, 0, dsts, offsets, 1); - REQUIRE(ret == -1); + // Poison the context with an error + ctx.error = HMLL_ERR(HMLL_ERR_IO_ERROR); - ctx.error = HMLL_OK; // reset so cleanup works - hmll_free_registry(®istry); - hmll_destroy(&ctx); - hmll_source_close(&src); + 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]") { @@ -648,7 +664,7 @@ TEST_CASE("fetchv - multiple large tensors interleaved", "[fetchv][safetensors]" for (const auto& [name, backend] : kBackends) { INFO("Backend: " << name); - REQUIRE_FALSE(hmll_check(hmll_loader_init(&ctx, &src, 1, HMLL_DEVICE_CPU, backend))); + 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"}; @@ -659,7 +675,7 @@ TEST_CASE("fetchv - multiple large tensors interleaved", "[fetchv][safetensors]" 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); + dsts[i] = hmll_get_buffer_for_range(&ctx, r); offsets[i] = r.start; total += dsts[i].size; REQUIRE(dsts[i].size > 512 * 1024u); @@ -682,68 +698,76 @@ 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))); + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); - 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); + 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(), backend))); - 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; - } + 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); - 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); + 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, r); + offsets[i] = r.start; + } - for (size_t i = 0; i < N; ++i) hmll_free_buffer(&dsts[i]); - hmll_free_registry(®istry); - hmll_destroy(&ctx); - hmll_source_close(&src); + 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); + for (const auto& [name, backend] : kBackends) { + INFO("Backend: " << name); + + 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(), backend))); + + 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, r), + {.size = 0, .ptr = nullptr, .device = hmll_device_cpu()}, // zero-size + hmll_get_buffer_for_range(&ctx, 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 diff --git a/tests/tests_hmll_integration_safetensors.cpp b/tests/tests_hmll_integration_safetensors.cpp index 65d8a61..2c08ee7 100644 --- a/tests/tests_hmll_integration_safetensors.cpp +++ b/tests/tests_hmll_integration_safetensors.cpp @@ -70,7 +70,7 @@ TEST_CASE("safetensors integration - read multi-dtype file", "[safetensors][inte constexpr std::array backends = { std::make_pair("MMAP", HMLL_FETCHER_MMAP) }; #endif - for (auto [name, backend] : backends) { + for (const auto [name, backend] : backends) { INFO("Testing with backend: " << name); SECTION("can open and parse safetensors file") { diff --git a/tests/tests_hmll_safetensors.cpp b/tests/tests_hmll_safetensors.cpp index d3ab79f..83fcfca 100644 --- a/tests/tests_hmll_safetensors.cpp +++ b/tests/tests_hmll_safetensors.cpp @@ -65,9 +65,9 @@ TEST_CASE("safetensors dtype parsing", "[safetensors]") TEST_CASE("safetensors path creation", "[safetensors]") { SECTION("create valid path") { - const char* base = "/path/to/model.safetensors"; - const char* file = "model-00002-of-00005.safetensors"; - char* result = hmll_safetensors_path_create(base, file); + const auto base = "/path/to/model.safetensors"; + const auto file = "model-00002-of-00005.safetensors"; + auto result = hmll_safetensors_path_create(base, file); REQUIRE(result != nullptr); REQUIRE(strcmp(result, "/path/to/model-00002-of-00005.safetensors") == 0); @@ -81,9 +81,9 @@ TEST_CASE("safetensors path creation", "[safetensors]") } SECTION("handle path without directory") { - const char* base = "model.safetensors"; - const char* file = "other.safetensors"; - char* result = hmll_safetensors_path_create(base, file); + const auto base = "model.safetensors"; + const auto file = "other.safetensors"; + auto result = hmll_safetensors_path_create(base, file); REQUIRE(result != nullptr); REQUIRE(strcmp(result, "other.safetensors") == 0); @@ -95,9 +95,9 @@ TEST_CASE("safetensors path creation", "[safetensors]") TEST_CASE("safetensors offset parsing", "[safetensors]") { SECTION("parse valid offsets") { - const char* json = "[100, 500]"; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "[100, 500]"; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_offsets(root, &tensor); @@ -110,9 +110,9 @@ TEST_CASE("safetensors offset parsing", "[safetensors]") } SECTION("parse offsets with large values") { - const char* json = "[1073741824, 2147483648]"; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "[1073741824, 2147483648]"; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_offsets(root, &tensor); @@ -125,9 +125,9 @@ TEST_CASE("safetensors offset parsing", "[safetensors]") } SECTION("reject malformed offsets - not array") { - const char* json = "\"not an array\""; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "\"not an array\""; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_offsets(root, &tensor); @@ -139,9 +139,9 @@ TEST_CASE("safetensors offset parsing", "[safetensors]") } SECTION("reject malformed offsets - single element") { - const char* json = "[100]"; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "[100]"; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_offsets(root, &tensor); @@ -153,9 +153,9 @@ TEST_CASE("safetensors offset parsing", "[safetensors]") } SECTION("reject malformed offsets - empty array") { - const char* json = "[]"; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "[]"; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_offsets(root, &tensor); @@ -170,9 +170,9 @@ TEST_CASE("safetensors offset parsing", "[safetensors]") TEST_CASE("safetensors shape parsing", "[safetensors]") { SECTION("parse 1D shape") { - const char* json = "[1024]"; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "[1024]"; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_shape(root, &tensor); @@ -185,9 +185,9 @@ TEST_CASE("safetensors shape parsing", "[safetensors]") } SECTION("parse 2D shape") { - const char* json = "[32, 64]"; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "[32, 64]"; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_shape(root, &tensor); @@ -201,9 +201,9 @@ TEST_CASE("safetensors shape parsing", "[safetensors]") } SECTION("parse 4D shape") { - const char* json = "[2, 3, 224, 224]"; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "[2, 3, 224, 224]"; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_shape(root, &tensor); @@ -219,9 +219,9 @@ TEST_CASE("safetensors shape parsing", "[safetensors]") } SECTION("parse scalar (empty shape)") { - const char* json = "[]"; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "[]"; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_shape(root, &tensor); @@ -233,9 +233,9 @@ TEST_CASE("safetensors shape parsing", "[safetensors]") } SECTION("reject malformed shape - not array") { - const char* json = "42"; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "42"; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_shape(root, &tensor); @@ -247,9 +247,9 @@ TEST_CASE("safetensors shape parsing", "[safetensors]") } SECTION("parse large shape values") { - const char* json = "[4096, 4096]"; - yyjson_doc* doc = yyjson_read(json, strlen(json), 0); - yyjson_val* root = yyjson_doc_get_root(doc); + const auto json = "[4096, 4096]"; + auto doc = yyjson_read(json, strlen(json), 0); + auto root = yyjson_doc_get_root(doc); hmll_tensor_specs_t tensor = {}; hmll_error_t err = hmll_safetensors_header_parse_shape(root, &tensor);