From 2f76c2b4f4fcd3d6993c116bcf5936464c8e1b92 Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Mon, 14 Sep 2026 19:28:42 +0000 Subject: [PATCH 1/7] feat: add CPU decoding pipeline --- .github/workflows/cpu.yml | 40 ++ .github/workflows/tsan.yml | 4 - CMakeLists.txt | 22 +- CMakePresets.json | 10 +- Dockerfile | 6 +- README.md | 98 ++- bench/CMakeLists.txt | 16 +- bench/main.c | 194 ++++-- bench/scenarios/throughput-cpu.json | 55 ++ cmake/Helpers.cmake | 2 +- dev/cpu-pipeline-validation.md | 120 ++++ dev/cpu-pipeline.md | 178 ++++++ docs/api.md | 32 + docs/index.md | 31 +- docs/pipeline.md | 214 +++++++ docs/prefetch.md | 20 +- docs/troubleshooting.md | 18 +- flake.nix | 2 + mkdocs.yml | 3 +- pyproject.toml | 2 +- python/CMakeLists.txt | 34 +- python/damacy/__init__.py | 501 ++++++++++++--- python/damacy/_api.c | 146 +++-- python/damacy/_api.h | 27 +- python/damacy/_components.c | 320 ++++++++++ python/damacy/_native.c | 20 +- python/damacy/_native.pyi | 78 ++- python/tests/test_components.py | 376 +++++++++++ python/tests/test_damacy.py | 55 +- python/tests/test_deferred_release.py | 2 + src/CMakeLists.txt | 93 ++- src/assemble/assemble.h | 2 +- src/batch_pool/batch_pool.c | 3 +- src/damacy.h | 61 +- src/damacy_internal.h | 138 +--- src/damacy_lifecycle.c | 780 +++++++---------------- src/damacy_pipeline.h | 112 ++++ src/damacy_plan.c | 343 +--------- src/damacy_pop.c | 279 ++------ src/damacy_push.c | 82 +-- src/damacy_scheduler.c | 106 ++- src/decoder/blosc1_parse.cu | 17 +- src/{planner => executor}/coalesce.c | 6 +- src/{planner => executor}/coalesce.h | 10 +- src/executor/cpu_executor.c | 646 +++++++++++++++++++ src/executor/cuda_executor.c | 623 ++++++++++++++++++ src/executor/cuda_executor.h | 10 + src/executor/cuda_geometry.inc | 138 ++++ src/executor/cuda_stub.c | 23 + src/executor/dispatch.c | 140 ++++ src/executor/dispatch.h | 123 ++++ src/{planner => executor}/group_chunks.c | 8 +- src/{planner => executor}/group_chunks.h | 6 +- src/{planner => executor}/read_op_sort.c | 4 +- src/{planner => executor}/read_op_sort.h | 0 src/numa/affinity.c | 55 ++ src/numa/numa.c | 53 +- src/numa/numa.h | 9 - src/numa/numa_cuda.h | 10 + src/pipeline/components.c | 232 +++++++ src/pipeline/components.h | 113 ++++ src/pipeline/zarr_planner.c | 320 ++++++++++ src/planner/plan.c | 12 + src/planner/plan.h | 60 ++ src/planner/plan_builder.c | 316 +++++++++ src/planner/plan_builder.h | 25 + src/planner/planner.c | 741 ++++----------------- src/planner/planner.h | 176 +---- src/prefetch/prefetcher.c | 15 +- src/render_job/render_job.c | 31 +- src/render_job/render_job.h | 10 +- src/scheduler/scheduler.c | 34 +- src/scheduler/scheduler.h | 10 +- src/threadpool/threadpool.c | 4 +- src/wave/wave_budget.c | 24 +- src/wave/wave_budget.h | 1 + src/wave/wave_pool.c | 2 +- tests/CMakeLists.txt | 11 +- tests/test_assemble.c | 2 +- tests/test_coalesce.c | 35 +- tests/test_cpu_pipeline.c | 354 ++++++++++ tests/test_damacy.c | 4 +- tests/test_damacy_blosc.c | 9 +- tests/test_damacy_plan.c | 151 ++++- tests/test_group_chunks.c | 45 +- tests/test_planner.c | 40 +- tests/test_render_job.c | 8 +- tests/test_scheduler.c | 14 +- tests/test_wave_pool.c | 2 +- tests/write_zarr.py | 9 +- 90 files changed, 6571 insertions(+), 2745 deletions(-) create mode 100644 .github/workflows/cpu.yml create mode 100644 bench/scenarios/throughput-cpu.json create mode 100644 dev/cpu-pipeline-validation.md create mode 100644 dev/cpu-pipeline.md create mode 100644 docs/pipeline.md create mode 100644 python/damacy/_components.c create mode 100644 python/tests/test_components.py create mode 100644 src/damacy_pipeline.h rename src/{planner => executor}/coalesce.c (97%) rename src/{planner => executor}/coalesce.h (82%) create mode 100644 src/executor/cpu_executor.c create mode 100644 src/executor/cuda_executor.c create mode 100644 src/executor/cuda_executor.h create mode 100644 src/executor/cuda_geometry.inc create mode 100644 src/executor/cuda_stub.c create mode 100644 src/executor/dispatch.c create mode 100644 src/executor/dispatch.h rename src/{planner => executor}/group_chunks.c (91%) rename src/{planner => executor}/group_chunks.h (86%) rename src/{planner => executor}/read_op_sort.c (96%) rename src/{planner => executor}/read_op_sort.h (100%) create mode 100644 src/numa/affinity.c create mode 100644 src/numa/numa_cuda.h create mode 100644 src/pipeline/components.c create mode 100644 src/pipeline/components.h create mode 100644 src/pipeline/zarr_planner.c create mode 100644 src/planner/plan.c create mode 100644 src/planner/plan.h create mode 100644 src/planner/plan_builder.c create mode 100644 src/planner/plan_builder.h create mode 100644 tests/test_cpu_pipeline.c diff --git a/.github/workflows/cpu.yml b/.github/workflows/cpu.yml new file mode 100644 index 00000000..20699e57 --- /dev/null +++ b/.github/workflows/cpu.yml @@ -0,0 +1,40 @@ +name: cpu + +on: + push: + branches: [main] + pull_request: + workflow_dispatch: + +jobs: + cpu: + runs-on: ubuntu-24.04 + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - name: Install dependencies + run: | + sudo apt-get update + sudo apt-get install -y cmake ninja-build pkg-config liburing-dev libzstd-dev libblosc-dev + python -m pip install uv pytest pytest-cov numpy + - name: Build without CUDA + run: | + cmake --preset cpu -DDAMACY_PYTHON=ON -DPython_EXECUTABLE="$(command -v python)" -DCMAKE_DISABLE_FIND_PACKAGE_CUDAToolkit=ON + cmake --build build -j2 + - name: Test + run: ctest --test-dir build --output-on-failure --timeout 180 -j2 + - name: Check native dependencies + env: + PYTHONPATH: build/python + run: | + python - <<'PY' + import subprocess + from damacy import _native + assert not _native.CUDA_ENABLED + dependencies = subprocess.check_output(['ldd', _native.__file__], text=True).lower() + print(dependencies) + assert all(name not in dependencies for name in ('libcuda', 'libcudart', 'libnvcomp')) + PY diff --git a/.github/workflows/tsan.yml b/.github/workflows/tsan.yml index d7b44529..6bce38ae 100644 --- a/.github/workflows/tsan.yml +++ b/.github/workflows/tsan.yml @@ -3,10 +3,6 @@ # unsynchronized access to shared state surfaces as a TSan report rather than # a flaky hang. # -# The test itself is pure C, but the test binary links transitively -# against libcuda via the shared test fixture / zarr / store libs, so -# it needs the NVIDIA driver visible at load time. Uses -# --device=nvidia.com/gpu=all. name: tsan diff --git a/CMakeLists.txt b/CMakeLists.txt index eff46d72..667bd807 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -15,11 +15,16 @@ include(Warnings) include(Helpers) include(Fuzz) # declares DAMACY_FUZZ and applies fuzz-mode flags +option(DAMACY_CUDA "Build the CUDA executor" ON) +if(DAMACY_FUZZ) + set(DAMACY_CUDA OFF) +endif() + find_package(Threads REQUIRED) find_package(PkgConfig REQUIRED) pkg_check_modules(LIBURING REQUIRED IMPORTED_TARGET liburing) -if(NOT DAMACY_FUZZ) +if(DAMACY_CUDA) if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) set(CMAKE_CUDA_ARCHITECTURES "75-real;80-real;86-real;89-real;90-real;90-virtual" @@ -48,6 +53,12 @@ option( # inline no-ops via static inlines in src/nvtx/nvtx.h. Defaults ON for # non-fuzz builds; src/CMakeLists.txt forces OFF under DAMACY_FUZZ. option(DAMACY_NVTX_ENABLED "Enable NVTX timeline instrumentation" ON) +if(NOT DAMACY_CUDA) + set(DAMACY_NVTX_ENABLED OFF) + if(DAMACY_ENABLE_GDS) + message(FATAL_ERROR "DAMACY_ENABLE_GDS requires DAMACY_CUDA") + endif() +endif() # DAMACY_ENABLE_GDS — gates the cuFile/GDS reader backend # (src/store/store_fs_gds.c). Defaults OFF so developers/CI without @@ -57,7 +68,7 @@ option(DAMACY_ENABLE_GDS "Build cuFile/GDS reader backend" OFF) # cufile.h is needed at compile time for typedefs; libcufile.so.0 is # dlopen'd at runtime (see src/store/store_fs_gds.c), not linked. Gated # on DAMACY_ENABLE_GDS so non-GDS builds don't require cufile.h. -if(NOT DAMACY_FUZZ AND DAMACY_ENABLE_GDS) +if(DAMACY_CUDA AND DAMACY_ENABLE_GDS) find_package(CuFile REQUIRED) endif() @@ -87,6 +98,11 @@ if(DAMACY_TSAN AND NOT DAMACY_FUZZ) target_link_options(warnings INTERFACE -fsanitize=thread) endif() +if(NOT DAMACY_FUZZ) + pkg_check_modules(ZSTD REQUIRED IMPORTED_TARGET libzstd) + pkg_check_modules(BLOSC REQUIRED IMPORTED_TARGET blosc) +endif() + add_subdirectory(src) add_subdirectory(bench) @@ -95,7 +111,7 @@ add_subdirectory(bench) # `cmake --build build --target asmbench` etc. The directory may be # pruned from build contexts that don't ship experiments (e.g. Docker). if( - NOT DAMACY_FUZZ + DAMACY_CUDA AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/experiments/asm/CMakeLists.txt" ) add_subdirectory(experiments/asm) diff --git a/CMakePresets.json b/CMakePresets.json index 8ad5067c..489ea369 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -7,7 +7,15 @@ "binaryDir": "${sourceDir}/build", "cacheVariables": { "CMAKE_EXPORT_COMPILE_COMMANDS": "ON", - "CMAKE_BUILD_TYPE": "RelWithDebInfo" + "CMAKE_BUILD_TYPE": "RelWithDebInfo", + "DAMACY_CUDA": "ON" + } + }, + { + "name": "cpu", + "inherits": "default", + "cacheVariables": { + "DAMACY_CUDA": "OFF" } }, { diff --git a/Dockerfile b/Dockerfile index 40c1e2ce..d16ca40c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -51,6 +51,8 @@ RUN apt-get update \ python3 \ python3-dev \ liburing-dev \ + libzstd-dev \ + libblosc-dev \ libnuma1 \ libmount1 \ libudev1 \ @@ -81,7 +83,7 @@ ENV PATH=/opt/venv/bin:${PATH} \ # pytest drives python/tests/* via the python_pytest ctest target; # installed alongside scikit-build-core so the cmake configure can # detect it and register the test. -RUN uv pip install scikit-build-core pytest pytest-cov +RUN uv pip install scikit-build-core pytest pytest-cov numpy # ----- build configuration --------------------------------------------------- # Override at `docker build` time via --build-arg to produce a coverage @@ -137,7 +139,7 @@ RUN ctest --test-dir build --output-on-failure -E "test_damacy|test_assemble|pyt # so the editable install resolves `damacy._native` without rebuilding. # Skipped under TSan (the .so isn't built; the install would 404). RUN if [ "${DAMACY_TSAN}" != "ON" ]; then \ - cp build/python/_native*.so python/damacy/ && \ + cp build/python/damacy/_native*.so python/damacy/ && \ uv pip install --no-deps --no-build-isolation -e .; \ fi diff --git a/README.md b/README.md index f14a6a55..bf9765b7 100644 --- a/README.md +++ b/README.md @@ -7,15 +7,67 @@ [![bench](https://github.com/nclack/damacy/actions/workflows/bench.yml/badge.svg)](https://nclack.github.io/damacy/throughput/) [![docs](https://github.com/nclack/damacy/actions/workflows/docs.yml/badge.svg)](https://nclack.github.io/damacy/) -High-speed streamed assembly of nD tensors from zarr sources to GPU. +Streamed assembly of n-dimensional tensors from Zarr sources into RAM or GPU memory. -Damacy reads sharded [NGFF](https://ngff.openmicroscopy.org/) [zarr -v3](https://zarr-specs.readthedocs.io/en/latest/v3/core/index.html) stores -directly on the GPU: per-shard chunk indexing, parallel host I/O, in-flight -GPU-side decompression (zstd, blosc1-zstd), and assembly of each batch -as a DLPack-ready device tensor. +Damacy loads array metadata and shard indexes, plans the chunks needed for a +batch, then reads, decodes, and assembles them with a CPU or CUDA executor. +Both paths support raw bytes, zstd, and Blosc-zstd and return contiguous, +DLPack-compatible tensors. CPU builds have no CUDA or nvCOMP dependency. -## Quick start +Source URIs identify concrete Zarr v3 arrays, including arrays inside +[NGFF](https://ngff.openmicroscopy.org/) multiscale images. Queries currently +select rectangular regions in array coordinates. Index queries, transformed +crops, and automatic NGFF level selection are planned extensions. + +## CPU quick start + +```python +import damacy +import numpy as np + +metadata_reader = damacy.FileMetadataReader(concurrency=64) +chunk_reader = damacy.FileReader(workers=8, max_inflight_reads=4096) +metadata = damacy.ZarrMetadata( + reader=metadata_reader, + cache=damacy.MetadataCache(array_entries=256, shard_index_entries=8192), +) +planner = damacy.ChunkPlanner(metadata=metadata, limits=damacy.PlanLimits()) +executor = damacy.CpuExecutor( + reader=chunk_reader, + limits=damacy.CpuLimits(max_memory_bytes=1 << 30, decode_workers=8), +) + +with damacy.Pipeline( + planner=planner, + executor=executor, + output=damacy.BatchSpec(samples=2, shape=(64, 256, 256), dtype="f32"), + queues=damacy.QueueLimits(lookahead_samples=4, prepared_batches=2), +) as pipeline: + pipeline.push([ + damacy.Sample(uri=uri, aabb=[(0, 64), (0, 256), (0, 256)]) + for uri in ["/data/image-1.zarr/0", "/data/image-2.zarr/0"] + ]) + with pipeline.pop() as batch: + array = np.from_dlpack(batch) + print(array.shape) + del array +``` + +The metadata provider loads descriptions on demand from the pushed URIs; +the chunk reader loads the encoded byte ranges that planning produces. +`np.from_dlpack` shares the output buffer. A live array keeps that buffer +occupied even after its `Batch` is released; use `.copy()` to retain data +while allowing the pool to reuse its storage. + +See [Pipeline composition](docs/pipeline.md) for the complete component, +limit, ownership, and build contracts. Build a CPU Python package from source +with the Linux dependencies `liburing`, `libzstd`, and `libblosc` installed: + +```sh +pip install . --config-settings=cmake.define.DAMACY_CUDA=OFF +``` + +## CUDA quick start ```python import random @@ -25,13 +77,9 @@ import torch cfg = damacy.Config( samples_per_batch=8, sample_shape=(64, 256, 256), - # Resource caps are fixed at construction; nothing grows after. max_gpu_memory_bytes=1 << 30, # primary GPU budget dtype="bf16", # source dtype is cast on assemble - # One pipeline binds to one GPU. Omit `device=` to capture the - # current CUDA context (handy single-GPU; PyTorch sets one up - # implicitly). For multi-GPU pass `device=local_rank` — see - # https://nclack.github.io/damacy/distributed/ + device=0, ) # A Sample names an absolute uri and a per-axis half-open AABB into @@ -99,7 +147,7 @@ Damacy reads zarr v3 (sharded and non-sharded). What's recognized today: | Array versions | v3 | v2 stores are not read | | Layout (sharded) | `sharding_indexed` | with `index_location` either `"start"` or `"end"` (default) | | Layout (non-sharded) | yes | each chunk is a separate file at `c///...` | -| Inner / chunk codec | `bytes` (passthrough), `zstd`, `blosc` (cname=`zstd`) | `blosc` with `cname=lz4`/`lz4hc` is recognized at parse time and rejected at planning | +| Inner / chunk codec | `bytes` (passthrough), `zstd`, `blosc` (cname=`zstd`) | `blosc` with `cname=lz4`/`lz4hc` is recognized at parse time and rejected by the executor | | Sharding index codec | `bytes` + `crc32c` | the shard index itself; the data codec is separate | | Missing chunks | yes — read as `fill_value` | per zarr v3 spec; sharded "empty" entries (`offset==nbytes==2^64−1`) and missing chunk files both route here | @@ -113,22 +161,22 @@ If you have data that uses one of the unsupported codecs and you'd like it added ## Runtime dependencies -Damacy links the core CUDA and Linux async-I/O pieces directly. Optional -features dlopen their backends lazily, so a damacy binary loads on any host -with the required core libraries; optional features simply turn off when their -library is not present. +All builds require Linux async metadata I/O and CPU codec libraries. CUDA +builds additionally link the NVIDIA driver and nvCOMP. Build with +`DAMACY_CUDA=OFF` to import and run on a host without a CUDA driver. -| Library | Required at runtime | What you lose if missing | How damacy finds it | -|---|---|---|---| -| `libcuda.so.1` | always | nothing — damacy cannot run without it | NVIDIA driver install (`/run/opengl-driver/lib`, `/usr/lib/x86_64-linux-gnu`, …) | -| `liburing.so` | always on Linux builds | nothing — async metadata stat/open/read/close uses io_uring | normal dynamic loader (`LD_LIBRARY_PATH`, `ld.so.cache`, rpath, …) | -| `libnuma.so.1` | optional | NUMA pinning of pinned-host slabs, bulk I/O workers, scheduler thread, and metadata io_uring driver thread (single-socket hosts: no effect) | `dlopen` via dynamic loader (`LD_LIBRARY_PATH`, `ld.so.cache`) | -| `libcufile.so.0` | optional | `damacy_config.enable_gds = DAMACY_GDS_ON` — direct read of compressed chunks into device memory via NVIDIA GPUDirect Storage | `dlopen` via dynamic loader; ships with the CUDA toolkit and with nvidia-fs. Requires a build with `-DDAMACY_ENABLE_GDS=ON` (default OFF) | -| `libmount.so.1`, `libudev.so.1` | required *if and only if* using GDS | cuFile dlopen's these at driver init even in compat mode | dynamic loader | +| Library | Used by | How it is loaded | +|---|---|---| +| `liburing` | CPU and CUDA: async metadata I/O | normal dynamic loader | +| `libzstd`, `libblosc` | CPU decoding, included in both builds | normal dynamic loader | +| `libcuda.so.1`, nvCOMP | CUDA builds | driver loader; nvCOMP may be linked statically | +| `libnuma.so.1` | Optional CUDA placement and host affinity | `dlopen`; absence disables placement | +| `libcufile.so.0` | Optional CUDA GPUDirect Storage | `dlopen`; requires `DAMACY_ENABLE_GDS=ON` | +| `libmount.so.1`, `libudev.so.1` | cuFile initialization when GDS is used | dynamic loader | Metadata reads require a Linux kernel with the io_uring operations damacy uses: `STATX`, `OPENAT2`, `READ`, and `CLOSE`. If the kernel does not advertise -those operations, `damacy_create` fails instead of falling back to a thread +those operations, pipeline construction fails instead of falling back to a thread pool. GDS notes: diff --git a/bench/CMakeLists.txt b/bench/CMakeLists.txt index d194a0a4..5d068f70 100644 --- a/bench/CMakeLists.txt +++ b/bench/CMakeLists.txt @@ -10,7 +10,9 @@ target_link_libraries( PRIVATE damacy json json_writer slice strbuf warnings ) # nvcomp_static drags in libstdc++/libgcc_s; fold them into the binary too. -target_link_options(damacy_bench PRIVATE -static-libstdc++ -static-libgcc) +if(DAMACY_CUDA) + target_link_options(damacy_bench PRIVATE -static-libstdc++ -static-libgcc) +endif() add_executable(metadata_store_async_bench metadata_store_async_bench.c) target_link_libraries( @@ -19,8 +21,10 @@ target_link_libraries( ) # Mount-ceiling probe; replays DAMACY_TRACE_READS captures (see file header). -add_executable(preadreplay preadreplay.c) -target_link_libraries( - preadreplay - PRIVATE warnings Threads::Threads CUDA::cuda_driver -) +if(DAMACY_CUDA) + add_executable(preadreplay preadreplay.c) + target_link_libraries( + preadreplay + PRIVATE warnings Threads::Threads CUDA::cuda_driver + ) +endif() diff --git a/bench/main.c b/bench/main.c index 6750e8b4..0d3461d9 100644 --- a/bench/main.c +++ b/bench/main.c @@ -6,20 +6,23 @@ // // All path/timestamp orchestration belongs in bench/run.py; this binary // is the timing core only. -#include "damacy.h" +#include "damacy_pipeline.h" #include "util/json.h" #include "util/json_writer.h" #include "util/slice.h" #include "util/strbuf.h" +#ifdef DAMACY_HAS_CUDA #include -#include +#endif + #include #include #include #include #include +#include #include #define countof(a) (sizeof(a) / sizeof((a)[0])) @@ -114,7 +117,10 @@ struct scenario uint32_t lookahead_samples; uint32_t n_io_threads; uint32_t metadata_io_concurrency; - uint64_t max_gpu_memory_bytes; // required + uint64_t max_gpu_memory_bytes; + uint64_t max_cpu_memory_bytes; + uint32_t decode_workers; + int cpu; uint32_t max_chunk_uncompressed_bytes; // 0 → tuning_defaults() baseline uint64_t max_read_op_bytes; // 0 → tuning_defaults() baseline uint32_t n_array_meta_cache; @@ -430,6 +436,29 @@ parse_scenario(struct cslice src, struct scenario* sc) sc->n_io_threads = (uint32_t)v; read_uint_opt(src, p_meta_io, countof(p_meta_io), &v, 8); sc->metadata_io_concurrency = (uint32_t)v; + static const struct json_query p_executor[] = { + { QUERY_KEY, .key = "pipeline" }, { QUERY_KEY, .key = "executor" } + }; + char executor[16] = "cuda"; + read_string_into( + src, p_executor, countof(p_executor), executor, sizeof(executor)); + if (strcmp(executor, "cpu") && strcmp(executor, "cuda")) + return 1; + sc->cpu = !strcmp(executor, "cpu"); + static const struct json_query p_cpu[] = { { QUERY_KEY, .key = "pipeline" }, + { QUERY_KEY, + .key = "max_cpu_memory_mb" } }; + static const struct json_query p_workers[] = { + { QUERY_KEY, .key = "pipeline" }, { QUERY_KEY, .key = "decode_workers" } + }; + read_uint_opt(src, p_cpu, countof(p_cpu), &v, 0); + if (v > (UINT64_MAX >> 20) || (sc->cpu && !v)) + return 1; + sc->max_cpu_memory_bytes = v << 20; + read_uint_opt(src, p_workers, countof(p_workers), &v, 8); + if (!v || v > UINT32_MAX) + return 1; + sc->decode_workers = (uint32_t)v; read_uint_opt(src, p_g, countof(p_g), &v, 0); sc->max_gpu_memory_bytes = v << 20; read_uint_opt(src, p_c, countof(p_c), &v, 0); @@ -649,8 +678,12 @@ array_table_init_uris(struct array_table* t, if (e != JSON_OK || v.type != JSON_STRING) return 1; char* slot = &t->uris[(size_t)t->n * BENCH_MAX_URI]; - int w = snprintf( - slot, BENCH_MAX_URI, "%s/%.*s", store_root, (int)cslice_len(v.s), v.s.beg); + int w = snprintf(slot, + BENCH_MAX_URI, + "%s/%.*s", + store_root, + (int)cslice_len(v.s), + v.s.beg); if (w < 0 || w >= BENCH_MAX_URI) return 1; int64_t shp[DAMACY_MAX_RANK]; @@ -712,6 +745,7 @@ struct run_metrics uint64_t pushed; uint64_t popped; struct damacy_stats stats; + uint64_t peak_host_bytes; }; // Push exactly n_target_batches * samples_per_batch samples and pop @@ -902,10 +936,10 @@ emit_results(const struct scenario* sc, const struct run_metrics* rm, FILE* out) emit_metric(&jw, &rm->stats.plan, "batch"); emit_metric(&jw, &rm->stats.io, "wave"); emit_metric(&jw, &rm->stats.input_transfer, "wave"); - emit_metric(&jw, &rm->stats.decode, "wave"); + emit_metric(&jw, &rm->stats.decode, sc->cpu ? "chunk" : "wave"); emit_metric(&jw, &rm->stats.post_decode, "wave"); emit_metric(&jw, &rm->stats.decode_gap, "wave"); - emit_metric(&jw, &rm->stats.assemble, "wave"); + emit_metric(&jw, &rm->stats.assemble, sc->cpu ? "chunk" : "wave"); emit_metric(&jw, &rm->stats.bind_wait, "wave"); emit_metric(&jw, &rm->stats.pop_wait, "poll"); jw_array_end(&jw); @@ -980,6 +1014,10 @@ emit_results(const struct scenario* sc, const struct run_metrics* rm, FILE* out) emit_metadata_op_latency(&jw, rm); jw_key(&jw, "gpu_bytes_committed"); jw_uint(&jw, rm->stats.gpu_bytes_committed); + jw_key(&jw, "host_bytes_committed"); + jw_uint(&jw, rm->stats.host_bytes_committed); + jw_key(&jw, "peak_host_bytes"); + jw_uint(&jw, rm->peak_host_bytes); jw_object_end(&jw); // Derived numbers. @@ -1035,6 +1073,85 @@ emit_results(const struct scenario* sc, const struct run_metrics* rm, FILE* out) // ---- main ------------------------------------------------------------------- +struct benchmark_pipeline +{ + struct damacy* handle; + struct damacy_reader* reader; + struct damacy_metadata_reader* metadata_reader; + struct damacy_metadata* metadata; + struct damacy_planner* planner; + struct damacy_executor* executor; +}; + +static void +pipeline_destroy(struct benchmark_pipeline* pipeline) +{ + damacy_destroy(pipeline->handle); + damacy_executor_destroy(pipeline->executor); + damacy_planner_destroy(pipeline->planner); + damacy_metadata_destroy(pipeline->metadata); + damacy_metadata_reader_destroy(pipeline->metadata_reader); + damacy_reader_destroy(pipeline->reader); +} + +static enum damacy_status +pipeline_create(const struct scenario* scenario, + const struct damacy_config* cfg, + struct benchmark_pipeline* pipeline) +{ + if (!scenario->cpu) + return damacy_create(cfg, &pipeline->handle); + enum damacy_status status = damacy_file_reader_create( + cfg->tuning.n_io_threads, 4096, &pipeline->reader); + if (status != DAMACY_OK) + return status; + status = + damacy_file_metadata_reader_create(cfg->tuning.metadata_io_concurrency, + &cfg->debug.metadata_latency, + &pipeline->metadata_reader); + if (status != DAMACY_OK) + return status; + status = damacy_zarr_metadata_create( + pipeline->metadata_reader, + &(struct damacy_metadata_cache_config){ cfg->tuning.n_array_meta_cache, + cfg->tuning.n_shard_index_cache }, + &pipeline->metadata); + if (status != DAMACY_OK) + return status; + status = damacy_chunk_planner_create( + pipeline->metadata, + &(struct damacy_plan_limits){ + .max_chunks = 16384, + .max_chunk_bytes = cfg->tuning.max_chunk_uncompressed_bytes, + .max_shards_per_sample = cfg->tuning.max_shards_per_sample, + .max_plan_bytes = 64ull << 20 }, + &pipeline->planner); + if (status != DAMACY_OK) + return status; + status = damacy_cpu_executor_create( + pipeline->reader, + &(struct damacy_cpu_config){ + .decode_workers = scenario->decode_workers, + .max_encoded_chunk_bytes = (uint32_t)cfg->tuning.max_read_op_bytes, + .max_decoded_chunk_bytes = cfg->tuning.max_chunk_uncompressed_bytes, + .max_memory_bytes = scenario->max_cpu_memory_bytes }, + &pipeline->executor); + if (status != DAMACY_OK) + return status; + struct damacy_batch_spec output = { .dtype = cfg->dtype, + .sample_rank = cfg->sample_rank, + .samples_per_batch = + cfg->samples_per_batch }; + memcpy(output.sample_shape, cfg->sample_shape, sizeof(output.sample_shape)); + return damacy_pipeline_create( + pipeline->planner, + pipeline->executor, + &output, + &(struct damacy_queue_limits){ .lookahead_samples = cfg->lookahead_samples, + .prepared_batches = 2 }, + &pipeline->handle); +} + int main(int argc, char** argv) { @@ -1082,13 +1199,14 @@ main(int argc, char** argv) } } - fprintf(stderr, - "scenario: store_root=%s n_arrays=%u rank=%u batches=%u (warmup=%u)\n", - sc.store_root, - at.n, - sc.rank, - sc.n_batches, - sc.n_warmup_batches); + fprintf( + stderr, + "scenario: store_root=%s n_arrays=%u rank=%u batches=%u (warmup=%u)\n", + sc.store_root, + at.n, + sc.rank, + sc.n_batches, + sc.n_warmup_batches); struct damacy_config cfg = { .samples_per_batch = sc.samples_per_batch, @@ -1120,30 +1238,29 @@ main(int argc, char** argv) struct rng rng = { .s = sc.sampling_seed ? sc.sampling_seed : 0xdeadbeefULL }; - // damacy_create requires a CUcontext current. Retain dev 0's primary. - if (cuInit(0) != CUDA_SUCCESS) { - fprintf(stderr, "cuInit failed\n"); - array_table_free(&at); - free(json_buf); - return 1; - } - CUdevice cu_dev = 0; - CUcontext cu_ctx = NULL; - if (cuDeviceGet(&cu_dev, 0) != CUDA_SUCCESS || - cuDevicePrimaryCtxRetain(&cu_ctx, cu_dev) != CUDA_SUCCESS || - cuCtxSetCurrent(cu_ctx) != CUDA_SUCCESS) { - fprintf(stderr, "primary ctx setup failed\n"); - array_table_free(&at); - free(json_buf); - return 1; +#ifdef DAMACY_HAS_CUDA + if (!sc.cpu) { + CUdevice device; + CUcontext context; + if (cuInit(0) != CUDA_SUCCESS || cuDeviceGet(&device, 0) != CUDA_SUCCESS || + cuDevicePrimaryCtxRetain(&context, device) != CUDA_SUCCESS || + cuCtxSetCurrent(context) != CUDA_SUCCESS) { + fprintf(stderr, "CUDA context initialization failed\n"); + array_table_free(&at); + free(json_buf); + return 1; + } } +#endif double t_init_a = now_seconds(); - struct damacy* d = NULL; - enum damacy_status cs = damacy_create(&cfg, &d); + struct benchmark_pipeline pipeline = { 0 }; + enum damacy_status cs = pipeline_create(&sc, &cfg, &pipeline); + struct damacy* d = pipeline.handle; double t_init_b = now_seconds(); if (cs != DAMACY_OK) { - fprintf(stderr, "damacy_create: %s\n", damacy_status_str(cs)); + fprintf(stderr, "pipeline_create: %s\n", damacy_status_str(cs)); + pipeline_destroy(&pipeline); array_table_free(&at); free(json_buf); return 1; @@ -1172,7 +1289,7 @@ main(int argc, char** argv) NULL, NULL, NULL)) { - damacy_destroy(d); + pipeline_destroy(&pipeline); array_table_free(&at); free(json_buf); return 1; @@ -1201,7 +1318,7 @@ main(int argc, char** argv) &rm.consumer_block_ms_total, &rm.consumer_push_ms_total, &rm.consumer_pop_wait_ms_total)) { - damacy_destroy(d); + pipeline_destroy(&pipeline); array_table_free(&at); free(json_buf); return 1; @@ -1220,7 +1337,7 @@ main(int argc, char** argv) &rm.consumer_block_ms_total, &rm.consumer_push_ms_total, &rm.consumer_pop_wait_ms_total)) { - damacy_destroy(d); + pipeline_destroy(&pipeline); array_table_free(&at); free(json_buf); return 1; @@ -1233,7 +1350,10 @@ main(int argc, char** argv) rm.popped += popped_steady; damacy_stats_get(d, &rm.stats); - damacy_destroy(d); + struct rusage usage; + if (getrusage(RUSAGE_SELF, &usage) == 0) + rm.peak_host_bytes = (uint64_t)usage.ru_maxrss * 1024; + pipeline_destroy(&pipeline); array_table_free(&at); emit_results(&sc, &rm, stdout); diff --git a/bench/scenarios/throughput-cpu.json b/bench/scenarios/throughput-cpu.json new file mode 100644 index 00000000..e3c6e127 --- /dev/null +++ b/bench/scenarios/throughput-cpu.json @@ -0,0 +1,55 @@ +{ + "name": "throughput-cpu", + "dataset": { + "store_root": "/mnt/main0/home/nclack/data/damacy/throughput", + "n_zarrs": 16, + "uri_fmt": "z%03u/scale0/image", + "zarr_shape": [ + 512, + 512, + 512 + ], + "chunk_shape": [ + 32, + 64, + 64 + ], + "shard_shape": [ + 256, + 256, + 512 + ], + "dtypes": [ + "u16" + ], + "codecs": [ + "zstd" + ], + "clevel": 3, + "entropy": 0.5, + "seed": 42 + }, + "sampling": { + "sample_shape": [ + 128, + 128, + 128 + ], + "n_batches": 12, + "n_warmup_batches": 3, + "samples_per_batch": 128, + "seed": 1234 + }, + "pipeline": { + "dtype": "f32", + "lookahead_samples": 512, + "n_io_threads": 8, + "metadata_io_concurrency": 64, + "n_array_meta_cache": 8192, + "n_shard_index_cache": 32768, + "max_shards_per_sample": 8, + "executor": "cpu", + "max_cpu_memory_mb": 4096, + "decode_workers": 8 + } +} diff --git a/cmake/Helpers.cmake b/cmake/Helpers.cmake index 45c0ccbb..d8b19731 100644 --- a/cmake/Helpers.cmake +++ b/cmake/Helpers.cmake @@ -25,7 +25,7 @@ endfunction() # Same as add_src_lib but a no-op under DAMACY_FUZZ (fuzz mode doesn't # enable CUDA, so any CUDA-linking target would fail to configure). function(add_cuda_lib TARGET) - if(DAMACY_FUZZ) + if(NOT DAMACY_CUDA) return() endif() add_src_lib(${TARGET} ${ARGN}) diff --git a/dev/cpu-pipeline-validation.md b/dev/cpu-pipeline-validation.md new file mode 100644 index 00000000..27ec0e2e --- /dev/null +++ b/dev/cpu-pipeline-validation.md @@ -0,0 +1,120 @@ +# CPU pipeline validation — 2026-09-14 + +The CPU executor uses ordinary host memory and shares rectangular query +preparation with the CUDA executor. The construction and lifetime contract is +in [Pipeline composition](../docs/pipeline.md); the extension plan for index +queries and NGFF resampling is in [CPU pipeline and query architecture](cpu-pipeline.md). + +## Workload + +The existing throughput dataset contains 16 sharded Zarr v3 arrays of shape +`512 × 512 × 512`, with `u16` data and zstd compression. Chunks are +`32 × 64 × 64`; shards are `256 × 256 × 512`. Data generation uses seed 42; +crop selection uses seed 1234. Each batch contains 128 crops of shape +`128 × 128 × 128`, converted to `f32`: 1 GiB of useful output per batch. + +The CPU measurements use three warmup batches and twelve measured batches, +eight I/O workers, and a 4 GiB executor memory limit. Each worker count was +measured once in the same 16-core CPU allocation, in the order 1, 4, 8, 16. +Filesystem caches were shared between runs. These are rates after warmup, not +cold-storage measurements or training-loop timings. The scenario +is [throughput-cpu.json](../bench/scenarios/throughput-cpu.json). + +CUDA comparisons use five warmup batches and thirty measured batches, with a +6 GiB device-memory limit. Baseline and refactor runs alternate on one L40 +with eight allocated CPU cores. Both harnesses initialize the CUDA context +before creating the pipeline and use the legacy captured-context API. Both +builds use `RelWithDebInfo` with the same C/C++ and CUDA compilers. The baseline +is an independent source snapshot of commit +`dee3e660c58009a1ae22a3774c963e16c9e086db`, compiled for `sm_89` with +the same CUDA 13.1.2 toolkit and nvCOMP 5.2.0.10. The CPU build uses GCC 11.4, +zstd 1.5.7, and C-Blosc 1.21.6. Builds and CPU tests run on CPU compute nodes. + +## CPU results + +Throughput is useful output bytes divided by measured wall time; GB/s is +decimal. Peak RAM is the benchmark process's maximum resident set size. + +| Decode workers | Useful output GB/s | Wall time for 12 GiB | Peak RAM GiB | +| ---: | ---: | ---: | ---: | +| 1 | 0.594 | 21.689 s | 2.015 | +| 4 | 1.815 | 7.100 s | 2.016 | +| 8 | 2.730 | 4.721 s | 2.013 | +| 16 | 4.104 | 3.139 s | 2.019 | + +Each run planned 67,727 chunk uses and decoded 53,413 distinct chunks, avoiding +21.1% of repeated decodes within batches. The measured work read +9,439,401,760 encoded bytes, decoded 14,001,897,472 source bytes, and produced +12,884,901,888 output bytes. There is no persistent decoded-chunk cache. + +The two 1 GiB output buffers dominate RAM. The executor's reservation grows +from 2.016 GiB with one worker to 2.255 GiB with sixteen, including conservative +codec workspace allowances. That reservation excludes metadata, plans, reader +queues, thread stacks, and allocator overhead; it is not a process RSS limit. + +## CUDA comparison and remaining performance issue + +| Pair | Baseline GB/s | Refactor GB/s | +| ---: | ---: | ---: | +| 1 | 11.834 | 11.684 | +| 2 | 12.516 | 11.716 | +| 3 | 12.409 | 11.646 | +| Median | 12.409 | 11.684 | + +Median useful throughput is **5.8% lower** after the refactor. This is an +open performance regression on this workload; CUDA throughput parity has not +been established. The lower first baseline result is retained in the table. +Three runs do not characterize all filesystem or scheduling variability. + +The measured decode kernels take about 2.515 seconds in both versions. The +additional time appears between decode waves, with substantial gaps at batch +boundaries. A separate timing trace recorded approximately 45 ms of shared +planning and 44 ms of CUDA dispatch preparation across 35 batches including +warmup. These measurements narrow the follow-up to host preparation and +scheduling; they do not establish one root cause. + +Correctness and retained-result lifetimes pass on CUDA. Further work on CUDA +preparation and scheduling should preserve the owned-plan boundary and be +measured separately from the future query features. + +## Checks + +- CPU configuration: all 23 CTest targets pass. Its Python run passes 22 tests + and skips 93 CUDA or optional-PyTorch cases. +- CUDA configuration on L40: all 35 CTest targets pass, including all 115 + Python tests with PyTorch installed. +- AddressSanitizer and UndefinedBehaviorSanitizer: all six selected metadata, + planning, scheduling, and CPU pipeline targets pass. +- The CPU wheel builds and imports without CUDA. A real crop and a retained + NumPy DLPack view work with no CUDA, cudart, or nvCOMP libraries in the + extension dependency chain or loaded process libraries. +- Four parser fuzzers build and pass 100-run smoke checks each. The public + Python API and new component tests pass Pyright; changed Python files pass + Ruff. Documentation builds with MkDocs strict checking. + +Correctness checks use independent expected values across raw, zstd, and Blosc +zstd encodings, supported source dtypes, missing fills, and `f32`/`bf16` +outputs. They also exercise duplicate chunk use, corrupt inputs, cache +eviction, bounded readers, memory limits, FIFO order, component reuse, +shutdown during a blocked pop, and results retained after shutdown. + +The CPU/CUDA comparison exposed and fixed the CUDA decoder's handling of +Blosc's uncompressed `MEMCPYED` payload. That payload follows the 16-byte +header directly; it does not contain the compressed-block offset table. +The layout follows [C-Blosc 1.21.6](https://github.com/Blosc/c-blosc/blob/v1.21.6/blosc/blosc.c). + +GDS and NUMA placement were unavailable on the L40 test host. This run does +not validate those optional paths or index queries, spatial interpolation, +NGFF level selection, or end-to-end training throughput. + +## Artifacts + +Scripts, complete test logs, source snapshot, scenarios, and raw benchmark +JSON are retained under +`~/tmp/2026-09-14-171630-damacy-cpu/`. CPU results are +`cpu-result-{1,4,8,16}.json`; the final CUDA comparison is in +`baseline-final-{1,2,3}.json` and `cuda-final-{1,2,3}.json`. Key validation logs +are `async-build.log`, `async-test-bench.log`, `async-ctest.log`, and +`package-finish.log`. Temporary timing instrumentation was removed from the +source and final binaries. Dataset files remain under +`~/data/damacy/throughput/`. diff --git a/dev/cpu-pipeline.md b/dev/cpu-pipeline.md new file mode 100644 index 00000000..15fc73aa --- /dev/null +++ b/dev/cpu-pipeline.md @@ -0,0 +1,178 @@ +# CPU pipeline and query architecture + +The CPU milestone introduces a shared preparation stage and separate CPU and +CUDA executors. Both support the existing rectangular queries. A CPU build +returns results in ordinary RAM and has no CUDA dependency. Index queries, +spatial resampling, and NGFF level selection remain later features. + +The public construction and lifetime contracts are in +[Pipeline composition](../docs/pipeline.md). The C factory API is +[damacy_pipeline.h](../src/damacy_pipeline.h); Python exposes the same components. +`Pipeline(Config(...))` and `damacy_create` compose the CUDA implementation for +existing callers. + +## Implemented separation + +| Component | Responsibility | Implementation | +| --- | --- | --- | +| Metadata reader/provider | Zarr descriptions, shard indexes, asynchronous metadata I/O, cache capacities. | `FileMetadataReader`, `ZarrMetadata`; active caches in `pipeline/zarr_planner.c`. | +| Planner | Validate rectangular requests and publish owned source/result plans. | `ChunkPlanner`, `planner/plan_builder.c`. | +| Executor | Read encoded data, prepare codecs, decode, assemble, manage output storage. | `executor/cpu_executor.c`, `executor/cuda_executor.c`. | +| Pipeline | Bound preparation, preserve batch order, retry backpressure, publish failures, coordinate shutdown. | `damacy_plan.c`, `damacy_scheduler.c`, `damacy_pop.c`. | + +The application constructs the services. Samples supply concrete array URIs; +the metadata provider loads descriptions on demand. Metadata and bulk chunk +reads use separate I/O queues. `FileReader` creates its bulk I/O queue when +constructed; no dataset is read until work is submitted. Pipeline construction starts metadata preparation and +the executor. + +Planner and executor operation tables are private. `start` binds workload +geometry and resources. The planner's `next` returns one complete owned plan or +`AGAIN`. Executor `submit` accepts ownership only on `OK`; `AGAIN` leaves the +same plan and batch ID with the pipeline. The executor releases accepted plans +when execution completes or stops. `take` returns ready results in batch order. +Failures stop admission and remain visible to subsequent pops until shutdown. + +A bounded queue allows metadata preparation to get ahead of execution. Owned +plans release cache protection early, so queued work does not depend on live +metadata cache entries. The metadata capacity floor is now +`lookahead_samples + samples_per_batch`, multiplied by the per-sample shard cap +for the shard cache. The legacy configuration validator keeps its older floor. + +The core scheduler calls interfaces without CUDA knowledge. CUDA context +binding happens on executor thread entry and exit. CUDA staging, streams, +layout probes, wave packing, device upload, and mutable cursors live under +execution. The old internal planner entry points are thin compatibility +adapters over the shared plan builder and CUDA dispatch builder; they do not +contain another chunk-enumeration implementation. + +## Owned plans + +[planner/plan.h](../src/planner/plan.h) contains the private representation: + +| Record | Meaning | +| --- | --- | +| Source array | Concrete array identity and a copy of its source description. | +| Source chunk | Absolute chunk coordinate, exact encoded range, decoded size, missing-chunk flag. | +| Result region | Output sample and a logical copy operation with its source region. | +| Chunk use | A connection from one distinct source chunk to one result region. | +| Output specification | Batch shape and destination dtype. | + +All paths and descriptors needed for execution belong to the plan. It contains +no cache handles, CUDA pointers, streams, physical output-slot identifiers, +staging offsets, or dispatch cursors. The output geometry is independent of +the source region, though rectangular copies currently require equal extents. +Sources are assumed unchanged during use; copying metadata is not a filesystem +snapshot. + +The plan records and storage are a separate library from the metadata-dependent +plan builder. Executors link the records and storage without the builder. + +The CUDA executor builds its own [dispatch records](../src/executor/dispatch.h) +from this plan. Alignment and read coalescing happen there, after metadata +resolution. It preserves the existing GPU decode/assembly path. CUDA chunk +layout probes are backend preparation and never enter shared planning. + +The CPU executor deduplicates decoded source chunks within a batch using the +plan's chunk-use links. It handles raw bytes, zstd, and C-Blosc zstd with no, +byte, or bit shuffle. It copies clipped chunk intersections and casts the +supported source types to `f32` or `bf16`, including fill-only chunks. The CUDA +executor currently retains its per-use decode behavior; cross-sample GPU decode +reuse can be optimized separately. + +## Resources and ownership + +CPU I/O workers and decode workers are configured independently. There are two +input groups and two output buffers. Each input group fits the decoder-worker +and reader capacities. Decoding uses a bounded per-worker output buffer and +zstd context; memory admission also reserves conservative Blosc workspace. +The memory cap covers executor buffers, not metadata, reader queues, plan +storage, thread stacks, allocator overhead, or total process RSS. + +The result handle has its own reference count and owns a reference to its +buffer. The executor owns another buffer reference and reuses storage only +when no result/consumer holds it. DLPack deleters release C references without +using Python objects, so consumer destruction need not run under the GIL. +Views survive batch release and pipeline shutdown. Explicit CUDA devices retain +the primary context and completion stream until the last buffer is released. +Caller-owned CUDA contexts must outlive their views. + +Planners, executors, metadata providers, and metadata readers reject use by two +active pipelines. They can be reused after shutdown. Python retains dependencies; +C borrows them until shutdown and requires reverse-order destruction. Queue and +buffer saturation report retriable backpressure, not a storage error. Closing +a pipeline stops preparation, wakes blocked pops, joins execution, and releases +queued work before its dependencies can disappear. + +## Next: index queries + +An index query should accept ordered index vectors along selected dimensions, +with ordinary contiguous ranges along the others. Preserve caller order and +repeated indices. Multiple indexed dimensions should use a Cartesian product, +matching MATLAB's per-dimension indexing; this differs from paired-coordinate +point queries. A stencil can generate an index vector, so no separate stencil +query type is needed. + +Extend logical result operations with compact index vectors and group chunk +uses during preparation. Keep output shape independent of the source bounding +range. Bound owned index storage along with query/plan storage. CPU assembly can +start with gather operations; CUDA can select an appropriate gather kernel. +Do not create one planning record per output voxel. + +## Next: spatial resampling and NGFF + +A spatial query requests a fixed output tensor, such as a transformed training +crop. Its transform maps output voxel centers into physical/source space. The +query includes sampler settings: interpolation, antialiasing/filter choices, +and boundary handling. These describe the result and belong with the query, +not in executor configuration. + +Resolve an NGFF image before enumerating source chunks: + +1. Load axes, units, level shapes, and level coordinate transforms. +2. Combine the output-grid transform with the level transforms and inspect + the sampling footprint, including rotation and anisotropic scales. +3. Choose an appropriate level for the requested sampling scale and filtering + policy, with an explicit query-level override when needed. +4. Resolve the output-to-array transform for that level and expand source + coverage by the sampler's interpolation/filter support. +5. Publish the concrete array, resolved transform, sampler, and output geometry + in the plan. Execution must not independently choose another level. + +Exact array-index queries continue to address their specified array. NGFF +scale-dependent loading applies to spatial resampling requests. The current +`ZarrMetadata` implementation does not parse an NGFF group or choose levels; +the private interfaces leave room to add that resolution stage. + +Resampling needs output tiles with all contributing source chunks available. +Keep decoded data until dependent tiles finish, and give each output tile one +writer. The shared plan already separates chunks from their uses, but the +current chunk-at-a-time rectangular executor will need this additional mode. +Boundary extension and filter halos must use the selected level's coordinates; +voxel-center conventions and downsampling quality need explicit tests. + +These additions should extend logical operations and query resolution rather +than mix source discovery with codec execution. A general operation graph, +persistent decoded-chunk cache, more codecs, and ragged batching are separate +work items. + +## Validation + +Measured results and build evidence are recorded in +[CPU pipeline validation](cpu-pipeline-validation.md). The current L40 +comparison shows about 6% lower CUDA throughput; further scheduling work is +needed before claiming performance parity. + +The CPU milestone checks independent crop values across codecs and source +dtypes, missing fills, bfloat16 conversion, duplicate chunk use, corrupt input, +resource limits, cache eviction, FIFO order, backpressure, and shutdown with +retained views. CPU builds and imports are checked without CUDA, including the +native library dependency chain. CUDA checks exercise both composition and the +legacy adapter, including PyTorch DLPack consumption and deferred release. + +Builds and CPU tests run on CPU compute nodes. CUDA tests and paired throughput +runs use one L40. The baseline is a source snapshot taken before the refactor; +benchmark inputs, seeds, output dtype, crop shape, and measured batch counts +are held constant. CPU measurements report useful output, decoded volume, +worker scaling, and peak process RAM. They are not a cold-storage benchmark or +a TensorStore comparison. diff --git a/docs/api.md b/docs/api.md index 4357be26..8ee83da1 100644 --- a/docs/api.md +++ b/docs/api.md @@ -12,6 +12,34 @@ The full public surface of the `damacy` package. ::: damacy.Batch +## Components + +::: damacy.FileMetadataReader + +::: damacy.ZarrMetadata + +::: damacy.ChunkPlanner + +::: damacy.FileReader + +::: damacy.CpuExecutor + +::: damacy.CudaExecutor + +## Output and limits + +::: damacy.BatchSpec + +::: damacy.QueueLimits + +::: damacy.MetadataCache + +::: damacy.PlanLimits + +::: damacy.CpuLimits + +::: damacy.CudaLimits + ## Value types ::: damacy.BatchInfo @@ -22,6 +50,8 @@ The full public surface of the `damacy` package. ## Enums +::: damacy.DeviceType + ::: damacy.Dtype ::: damacy.Status @@ -36,6 +66,8 @@ The full public surface of the `damacy` package. ::: damacy.NotFound +::: damacy.DeviceType + ::: damacy.DtypeMismatch ::: damacy.RankMismatch diff --git a/docs/index.md b/docs/index.md index 536b48ab..c3b28cbb 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,11 +1,13 @@ # damacy -**High-speed streamed assembly of tensors from zarr sources to GPU.** +**Streamed assembly of tensors from Zarr sources into RAM or GPU memory.** -damacy reads sharded [NGFF](https://ngff.openmicroscopy.org/) zarr stores -straight onto the GPU: per-shard chunk indexing, parallel host I/O, -in-flight GPU-side decompression (zstd, blosc1-zstd), and a typed -assemble kernel that lands each batch as a DLPack-ready device tensor. +Damacy separates metadata and chunk planning from CPU or CUDA execution. Both +executors read Zarr v3 arrays, decode raw bytes, zstd, or Blosc-zstd, and return +contiguous batches through DLPack. CPU builds require no CUDA toolkit or driver. + +Start with [Pipeline composition](pipeline.md) for the CPU API, complete +component construction, resource limits, and buffer ownership. [![build](https://github.com/nclack/damacy/actions/workflows/build.yml/badge.svg)](https://github.com/nclack/damacy/actions/workflows/build.yml) [![test](https://github.com/nclack/damacy/actions/workflows/test.yml/badge.svg)](https://github.com/nclack/damacy/actions/workflows/test.yml) @@ -13,7 +15,7 @@ assemble kernel that lands each batch as a DLPack-ready device tensor. --- -## Quick start +## CUDA quick start ```python import damacy @@ -24,6 +26,7 @@ cfg = damacy.Config( sample_shape=(64, 256, 256), max_gpu_memory_bytes=1 << 30, dtype="bf16", + device=0, ) samples = [ damacy.Sample(uri="/data/cells/cell-1.zarr", aabb=[(0, 64), (0, 256), (0, 256)]), @@ -39,7 +42,7 @@ with damacy.Pipeline(cfg) as d: ``` By default the pipeline captures whatever CUDA context is current on -the calling thread; PyTorch sets one up implicitly, and bare-Python +the calling thread; callers can initialize it through PyTorch, and bare-Python users can call `damacy._native.cuda_init_primary()` once. For multi-GPU setups, see [Distributed](distributed.md) for the device binding model and a torchrun example. @@ -47,23 +50,22 @@ binding model and a torchrun example. ## Concepts You hand damacy a stream of `Sample`s; it returns a stream of -`Batch`es, each one a device tensor of shape +`Batch`es, each one a CPU or CUDA tensor of shape `(samples_per_batch, *sample_shape)`. - A **`Sample`** is one crop request: a zarr URI plus an `aabb` (axis-aligned bounding box) given as a list of `(start, stop)` tuples — one per spatial axis. Every `aabb` must produce the same - per-sample shape, and that shape is `Config.sample_shape`. + per-sample shape, and that shape is `BatchSpec.shape` (or `Config.sample_shape`). - A **`Pipeline`** is a streaming context. You `push` an iterable of samples (lazy generators are fine — and recommended for long runs) and call `pop()` to block for the next ready batch. -- A **`Batch`** is a DLPack-ready handle to a GPU-resident tensor. - Use it inside a `with` block so damacy can reclaim the slot when - you're done. +- A **`Batch`** is a DLPack-ready handle to a CPU or CUDA tensor. + Use it inside a `with` block and release consumer views so damacy can + reclaim the buffer when you are done. `samples_per_batch`, `sample_shape`, and `max_gpu_memory_bytes` are required -on `Config`; everything else has a sensible default. The assemble -kernel casts heterogeneous source dtypes +on `Config`; everything else has a sensible default. Assembly casts heterogeneous source dtypes (`u8`/`u16`/`i16`/`u32`/`i32`/`f16`/`f32`) to the configured destination `dtype` (`f32` or `bf16`) on the way out, so your zarrs do not need to match it. @@ -74,6 +76,7 @@ The published API lives entirely under the top-level `damacy` package. The native extension (`damacy._native`) is an implementation detail documented only via its `.pyi` stub. +- [Pipeline composition](pipeline.md) — CPU and CUDA components, builds, and lifetimes. - [API reference](api.md) — `Pipeline`, `Config`, `Sample`, `Batch`, the exception hierarchy, and the `Stats`/`Metric` value types. - [GPU memory budget](budget.md) — how to think about diff --git a/docs/pipeline.md b/docs/pipeline.md new file mode 100644 index 00000000..86183e9e --- /dev/null +++ b/docs/pipeline.md @@ -0,0 +1,214 @@ +# Pipeline composition + +A pipeline receives a planner, an executor, an output specification, and queue +limits. The planner owns metadata preparation; the executor owns bulk reads, +decoding, and output buffers. CPU and CUDA execution use the same rectangular +queries and prepared-plan contract. + +## Construct a CPU pipeline + +```python +import damacy +import numpy as np + +metadata_reader = damacy.FileMetadataReader(concurrency=64) +chunk_reader = damacy.FileReader(workers=8, max_inflight_reads=4096) +metadata = damacy.ZarrMetadata( + reader=metadata_reader, + cache=damacy.MetadataCache(array_entries=256, shard_index_entries=8192), +) +planner = damacy.ChunkPlanner( + metadata=metadata, + limits=damacy.PlanLimits( + max_chunks=16384, + max_chunk_bytes=2 << 20, + max_shards_per_sample=64, + max_plan_bytes=64 << 20, + ), +) +executor = damacy.CpuExecutor( + reader=chunk_reader, + limits=damacy.CpuLimits( + max_memory_bytes=1 << 30, + decode_workers=8, + max_encoded_chunk_bytes=4 << 20, + max_decoded_chunk_bytes=2 << 20, + ), +) +output = damacy.BatchSpec(samples=2, shape=(64, 256, 256), dtype="f32") +queues = damacy.QueueLimits(lookahead_samples=4, prepared_batches=2) + +with damacy.Pipeline( + planner=planner, executor=executor, output=output, queues=queues +) as pipeline: + pipeline.push([ + damacy.Sample(uri=uri, aabb=[(0, 64), (0, 256), (0, 256)]) + for uri in ["/data/image-1.zarr/0", "/data/image-2.zarr/0"] + ]) + with pipeline.pop() as batch: + array = np.from_dlpack(batch) + assert array.shape == (2, 64, 256, 256) + assert batch.info.device_type == damacy.DeviceType.CPU + del array +``` + +The paths must name little-endian numeric Zarr v3 **arrays**. An NGFF group is +not resolved to a level automatically. The examples assume decoded chunks no larger than 2 MiB. + +`FileMetadataReader` supplies small asynchronous metadata reads. `ZarrMetadata` +supplies Zarr interpretation and cache capacities; `MetadataCache` is a capacity +configuration. The planner fetches array descriptions and shard indexes on +demand across the URIs in the sample stream. It publishes plans that own their +paths, source descriptions, exact encoded ranges, and result operations. +Metadata caches can evict entries after preparation finishes. + +`FileReader` has a separate queue for bulk chunk data. The CPU executor uses +ordinary RAM for input, decoded chunks, codec workspace, and output. It decodes +each distinct source chunk once per batch, then copies its intersections into +all samples that use it. There is no persistent decoded-chunk cache. + +## Select CUDA execution + +In a CUDA-enabled build, replace the executor with: + +```python +executor = damacy.CudaExecutor( + reader=chunk_reader, + device=0, + limits=damacy.CudaLimits(max_gpu_memory_bytes=1 << 30), +) +``` + +The remaining pipeline arguments and samples stay the same. Consume the +result with `torch.from_dlpack(batch)`. `device=0` retains that device's primary +CUDA context; `device=None` captures the caller's current context, which must +already exist. An externally managed context must outlive retained results. + +CUDA execution owns read alignment/coalescing, Blosc layout probes, staging +buffers, streams, nvCOMP decoding, and GPU assembly. Those details do not occur +in prepared plans. The existing `Pipeline(Config(...))` constructor remains a +CUDA convenience adapter. It cannot select the CPU executor. + +The `numa_strategy` argument to `CudaExecutor` controls its pinned buffers and +execution worker. +The injected readers inherit the caller's host affinity when their workers +start. The legacy `Config.numa_strategy` also applies while constructing those +readers, preserving placement of the complete legacy pipeline. + +## Limits and backpressure + +Sizes are bytes, with positive explicit limits. Zero is invalid for these +capacities. Defaults come from the Python value objects. + +| Setting | Scope | +| --- | --- | +| `BatchSpec` | Samples per batch, per-sample output shape, and `f32` or `bf16` output dtype. | +| `QueueLimits.lookahead_samples` | Sample requests waiting for metadata/preparation; at least one full batch. | +| `QueueLimits.prepared_batches` | Complete owned plans waiting for execution. | +| `MetadataCache` | Number of array descriptions and shard indexes retained by preparation. | +| `PlanLimits.max_chunks` | Chunk uses per batch, including chunks used by multiple samples. | +| `PlanLimits.max_chunk_bytes` | Decoded source bytes per chunk, before output conversion. | +| `PlanLimits.max_shards_per_sample` | Maximum number of shard files touched by one sample. | +| `PlanLimits.max_plan_bytes` | Owned storage per prepared plan. | +| `CpuLimits.max_memory_bytes` | CPU executor's input, decoded, codec-workspace, and two output buffers. | +| `CpuLimits.decode_workers` | Total decoding/assembly workers, including the calling scheduler thread. | +| `FileReader.workers` | Bulk I/O workers, separate from decoding workers. | +| `FileReader.max_inflight_reads` | Bulk read capacity; execution respects this bound and retries saturation. | +| `CudaLimits` | GPU memory and execution geometry, plus the CUDA codec-layout cache capacity. | + +For injected pipelines, define `floor = queues.lookahead_samples + output.samples`. +The metadata cache requires at least `floor` array entries and +`floor * plan_limits.max_shards_per_sample` shard-index entries. These floors +cover pending samples and the batch being prepared. Queued plans own their +metadata and do not pin cache entries. The legacy `Config` adapter retains its +older, stricter cache validation. + +CPU memory admission includes a conservative allowance for Blosc scratch +storage. `Stats.host_bytes_committed` reports that reservation, which can exceed +the bytes actually touched. It excludes metadata caches, prepared plans, reader +queues, thread stacks, and allocator overhead: it is not a process RSS limit. +Queued plan storage is bounded separately by +`prepared_batches * max_plan_bytes`, with up to two accepted plans and one plan +being built in addition. Retaining results across repeated pipeline restarts +also retains their allocations outside the new executor's budget. + +Both executors have two output buffers. A batch or a DLPack consumer holds its +buffer until all references are released. Holding both buffers pauses output +production. Release views promptly, or copy the result when it must be retained +while subsequent batches continue. Increasing `prepared_batches` increases +preparation capacity, not the output pool. + +Only complete batches are emitted. Trailing requests that do not fill a batch +are not returned. Invalid sample geometry is rejected at push; metadata, codec, +and memory-limit failures normally surface at pop. A terminal execution error +requires closing the pipeline and constructing another one. + +## Lifetimes and interoperation + +Use a context manager or call `Pipeline.close()`. Python retains the injected +components and their dependencies. Planners, executors, metadata providers, +and metadata readers each serve one active pipeline; simultaneous reuse is +rejected. They may be reused after close. Closing stops pending work and wakes +blocked pops. Reuse creates fresh active caches and execution resources. + +`np.from_dlpack(batch)` produces a CPU view; `torch.from_dlpack(batch)` accepts +CPU or CUDA results. A DLPack view remains valid after releasing the `Batch` and +closing the pipeline. Releasing the Python batch does not force a live view's +buffer back into the output pool. NumPy does not support `bf16` through DLPack; +use `f32` or a consumer with bfloat16 support. + +CPU batches report DLPack device `(1, 0)` and reject a stream argument other than +`None`. CUDA batches report `(2, device_id)` and preserve the existing stream +handoff. `BatchInfo.data` is the address on either device; `device_ptr` remains +an alias for compatibility. CPU data is ready when `pop()` returns and has no +ready stream. Use the device type before interpreting a raw pointer. + +The C factories and configuration structures are declared in +`src/damacy_pipeline.h`; push/pop and batch functions remain in `src/damacy.h`. +C pipelines borrow components until `damacy_shutdown`. Destroy dependencies in +reverse construction order. `damacy_batch_retain` / `damacy_batch_release` +manage result lifetimes independently of a pipeline; release each acquired +reference once. After destroying a pipeline, use `damacy_batch_release` for +retained results rather than a function requiring the old pipeline pointer. +Operation tables and prepared-plan records are private implementation APIs. + +## Build without CUDA + +On Linux, install C/C++ build tools, CMake, Ninja, pkg-config, liburing, zstd, +and C-Blosc development packages. For example, Ubuntu packages are +`build-essential cmake ninja-build pkg-config python3-dev liburing-dev libzstd-dev libblosc-dev`. + +```sh +cmake --preset cpu +cmake --build build +ctest --test-dir build --output-on-failure +``` + +The examples use NumPy (`pip install numpy`). The fixture generator uses `uv`. +Enable `DAMACY_PYTHON=ON` and provide a Python 3.11+ interpreter with pytest, +pytest-cov, and NumPy to include Python tests. +Build a CPU wheel with: + +```sh +pip install . --config-settings=cmake.define.DAMACY_CUDA=OFF +``` + +The extension has no CUDA or nvCOMP dependency in this configuration. +`CudaExecutor` reports that CUDA support was not built. CPU builds retain the +Linux io_uring metadata requirements. CUDA support is enabled by default; +turning it on builds both executors and requires the CUDA toolkit, nvCOMP, and +a runtime NVIDIA driver. GDS requires a CUDA build. + +## Future queries + +The first milestone implements rectangular copy/cast queries. Index-array +queries can add ordered selections and repeated indices without a separate +stencil type. Spatial queries will describe a fixed output tensor, a transform +from output coordinates to source space, and sampler settings including +interpolation, antialiasing, and boundary handling. + +For NGFF images, resolving a spatial query will choose an appropriate source +level from the requested sampling scale before enumerating chunks. The plan +will record the chosen array and transform. Resampling also needs interpolation +halos and dependencies across chunks. These are later planner and executor +operations; they are not implemented by the current box query. diff --git a/docs/prefetch.md b/docs/prefetch.md index ef84de24..a6a5ddc0 100644 --- a/docs/prefetch.md +++ b/docs/prefetch.md @@ -1,6 +1,6 @@ # Async prefetch -The [quick start](index.md#quick-start) shows the synchronous read +The [quick start](index.md#cuda-quick-start) shows the synchronous read pattern: ```python @@ -9,16 +9,15 @@ with batch as t: ... # train step ``` -`x` is a zero-copy view onto damacy's slot, and the `with` block -releases the slot at scope exit by host-syncing on damacy's -producer stream. That's fine when fwd/bwd happens inside the -block. +`x` shares the output buffer. Exiting the `with` block releases the batch +reference; the tensor retains its own reference until it is destroyed. Retaining +both output buffers prevents the next batch from completing. -Training loops that prefetch the next batch on a background thread -while the main thread runs fwd/bwd need to skip that host sync. -Both patterns below do — they hand a CUDA event (or stream) to -`Batch.release(event=...)`, and damacy waits on it before reusing -the slot. +For CUDA work that outlives the batch's host scope, pass a consumer event or +stream to `Batch.release(event=...)`. Damacy orders subsequent writes behind +that event. Release the tensor view as well when it is no longer needed. +These CUDA examples use an existing `Config`-based or injected pipeline; +[Pipeline composition](pipeline.md) covers CPU consumers. ## Deferred release (preferred) @@ -36,6 +35,7 @@ def prefetch(p): tensor, batch = prefetch_future.result() ... # forward / backward on tensor batch.release(event=torch.cuda.current_stream()) # no host sync +del tensor ``` `event=torch.cuda.current_stream()` records a CUevent on the diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 15ea55ce..2db80869 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -2,6 +2,7 @@ Common errors and what to check first. For multi-GPU specifics, see [Distributed → Common failures](distributed.md#common-failures). +CPU construction and limits are described in [Pipeline composition](pipeline.md). ## `InvalidArgument: no CUcontext is current` @@ -14,18 +15,16 @@ the calling thread yet. Two fixes: - Or prime a context implicitly before constructing the pipeline: `torch.empty(1, device="cuda")` is enough. -## `InvalidArgument` at `Pipeline(cfg)` from metadata I/O setup +## Pipeline construction fails during metadata I/O setup -The Linux metadata path uses io_uring for zarr metadata, shard-index, and -chunk-layout reads. At construction damacy requires kernel support for +The CPU and CUDA metadata paths use io_uring for Zarr metadata and shard indexes. At construction damacy requires kernel support for `IORING_OP_STATX`, `IORING_OP_OPENAT2`, `IORING_OP_READ`, and `IORING_OP_CLOSE`. If ring creation or the operation probe fails, -`Pipeline(cfg)` raises `InvalidArgument` rather than falling back to a legacy -thread pool. Check the native log for the exact io_uring failure. +pipeline construction fails without a thread-pool fallback. Check the native log for the exact io_uring failure. On supported kernels, an unusually high `metadata_io_concurrency` can also stress process file-descriptor limits because each in-flight metadata read can -hold an open fd. The default is 32; for much deeper settings, check +hold an open fd. The default is 64; for much deeper settings, check `ulimit -n` and remember to multiply by ranks per node. ## `BudgetExceeded` at `Pipeline(cfg)` @@ -39,9 +38,10 @@ the usual answer is to raise the cap. ## `BudgetExceeded` mid-stream A chunk's actual uncompressed size exceeds -`Config.max_chunk_uncompressed_bytes` (default 512 KiB). Raise +`PlanLimits.max_chunk_bytes` or `Config.max_chunk_uncompressed_bytes` (default 2 MiB). Raise that cap to fit the dataset, and raise `max_gpu_memory_bytes` -along with it if needed. See +along with it if needed. CPU executors also enforce `CpuLimits` for encoded +chunks, decoded chunks, and total executor buffers. See [GPU memory budget](budget.md#when-the-budget-refuses). ## `NotFound` or `DtypeMismatch` from `pop()` (not `push()`) @@ -72,7 +72,7 @@ The pool was empty for longer than `Config.pop_timeout_s` from previous batches — for example, stashing them in a list — which prevents damacy from reusing the underlying slot. Drop the references before the next `pop()`, or `.clone()` if you -genuinely need to keep them. +need to keep them. For a NumPy CPU view, use `.copy()`. ## `pop()` blocks forever diff --git a/flake.nix b/flake.nix index 22f3b50b..6b8c24b8 100644 --- a/flake.nix +++ b/flake.nix @@ -108,6 +108,8 @@ man-pages man-pages-posix liburing + zstd + c-blosc ]); CUDA_PATH = "${cudaPkgs.cudatoolkit}"; diff --git a/mkdocs.yml b/mkdocs.yml index f02c45e5..9461ea66 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,5 +1,5 @@ site_name: damacy -site_description: High-speed streamed assembly of tensors from zarr sources to GPU. +site_description: Streamed assembly of tensors from Zarr sources into RAM or GPU memory. site_url: https://nclack.github.io/damacy/ repo_url: https://github.com/nclack/damacy repo_name: nclack/damacy @@ -77,6 +77,7 @@ markdown_extensions: nav: - Home: index.md + - Pipeline composition: pipeline.md - GPU memory budget: budget.md - Distributed: distributed.md - Async prefetch: prefetch.md diff --git a/pyproject.toml b/pyproject.toml index 9e7d6f6b..8f3827a2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ authors = [{ name = "Nathan Clack", email = "nclack@gmail.com" }] requires-python = ">=3.11" [project.optional-dependencies] -dev = ["pytest>=8", "pytest-cov>=5", "ruff>=0.6", "pyright>=1.1"] +dev = ["pytest>=8", "pytest-cov>=5", "numpy>=1.26", "ruff>=0.6", "pyright>=1.1"] gpu = ["torch"] docs = [ "mkdocs>=1.6", diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 9656e2be..5aec7972 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -11,6 +11,7 @@ python_add_library( damacy/_native.c damacy/_log_sink.c damacy/_api.c + damacy/_components.c ) target_link_libraries(_native PRIVATE damacy log warnings) @@ -20,29 +21,44 @@ target_link_libraries(_native PRIVATE damacy log warnings) # wheel; the extension goes alongside the package's __init__.py. install(TARGETS _native LIBRARY DESTINATION damacy) -# --cov-fail-under floor matches codecov.yml's patch target so the two -# gates (project-total here, per-PR-patch on codecov) read the same number. +set_target_properties( + _native + PROPERTIES LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/damacy" +) +configure_file(damacy/__init__.py damacy/__init__.py COPYONLY) + if(BUILD_TESTING) - find_program(UV_EXECUTABLE uv) - if(UV_EXECUTABLE) + execute_process( + COMMAND "${Python_EXECUTABLE}" -c "import pytest, pytest_cov" + RESULT_VARIABLE PYTEST_AVAILABLE + OUTPUT_QUIET + ERROR_QUIET + ) + if(PYTEST_AVAILABLE EQUAL 0) + set(PYTHON_COVERAGE_ARGS + --cov=damacy + --cov-report=term-missing:skip-covered + ) + if(DAMACY_CUDA) + list(APPEND PYTHON_COVERAGE_ARGS --cov-fail-under=70) + endif() add_test( NAME python_pytest COMMAND - ${UV_EXECUTABLE} run --extra dev pytest -v --color=yes - --cov=damacy --cov-report=term-missing:skip-covered - --cov-fail-under=70 "${CMAKE_CURRENT_SOURCE_DIR}/tests" + "${Python_EXECUTABLE}" -m pytest -v ${PYTHON_COVERAGE_ARGS} + "${CMAKE_CURRENT_SOURCE_DIR}/tests" WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} ) set_tests_properties( python_pytest PROPERTIES ENVIRONMENT - "WRITE_ZARR_SCRIPT=${PROJECT_SOURCE_DIR}/tests/write_zarr.py" + "WRITE_ZARR_SCRIPT=${PROJECT_SOURCE_DIR}/tests/write_zarr.py;PYTHONPATH=${CMAKE_CURRENT_BINARY_DIR}" ) else() message( STATUS - "damacy: uv not on PATH; skipping python_pytest registration" + "damacy: install pytest and pytest-cov to enable python_pytest" ) endif() endif() diff --git a/python/damacy/__init__.py b/python/damacy/__init__.py index 6f3fe3ef..848992f5 100644 --- a/python/damacy/__init__.py +++ b/python/damacy/__init__.py @@ -37,6 +37,7 @@ import itertools import logging import math +import operator import os import threading import warnings @@ -55,21 +56,33 @@ __all__ = [ "Batch", "BatchInfo", + "BatchSpec", "BudgetExceeded", + "ChunkPlanner", "Config", + "CpuExecutor", + "CpuLimits", + "CudaExecutor", + "CudaLimits", "DamacyError", "DecodeError", + "DeviceType", "Dtype", "DtypeMismatch", + "FileMetadataReader", + "FileReader", "InvalidArgument", "LatencyModel", + "MetadataCache", "Metric", "NativeCudaError", "NotFound", "NumaStrategy", "OutOfMemory", "Pipeline", + "PlanLimits", "PoolStarved", + "QueueLimits", "RankMismatch", "Sample", "ShutdownError", @@ -77,6 +90,7 @@ "Status", "StorageError", "TryAgain", + "ZarrMetadata", "max_concurrency", "set_log_level", "set_log_quiet", @@ -94,8 +108,7 @@ def max_concurrency() -> int: class Dtype(IntEnum): - """Destination dtype for assembled batches. Sources may differ; the - assemble kernel casts each element to this type.""" + """Destination dtype for assembled batches. Sources are cast to this type.""" F32 = _native.DTYPE_F32 BF16 = _native.DTYPE_BF16 @@ -722,15 +735,278 @@ def default( ) +def _positive_int(value: int, name: str, maximum: int = (1 << 32) - 1) -> int: + value = operator.index(value) + if not 1 <= value <= maximum: + raise ValueError(f"{name} must be in [1, {maximum}]") + return value + + +def _component(factory: Any, *args: Any) -> Any: + try: + return factory(*args) + except _native.DamacyError as exc: + _reraise_typed(exc) + + +class DeviceType(IntEnum): + """DLPack memory location for a batch.""" + + CPU = 1 + CUDA = 2 + + +@dataclass(frozen=True, slots=True, init=False) +class BatchSpec: + """Fixed output tensor geometry, independent of the source metadata.""" + + samples: int + shape: tuple[int, ...] + dtype: Dtype = Dtype.F32 + + def __init__( + self, samples: int, shape: Sequence[int], dtype: Dtype | str | int = Dtype.F32 + ) -> None: + object.__setattr__(self, "samples", _positive_int(samples, "samples", 65535)) + dimensions = tuple( + _positive_int(x, "shape extent", (1 << 63) - 1) for x in shape + ) + if not 1 <= len(dimensions) <= 31: + raise ValueError("shape must have between 1 and 31 dimensions") + object.__setattr__(self, "shape", dimensions) + object.__setattr__(self, "dtype", Dtype.coerce(dtype)) + + +@dataclass(frozen=True, slots=True) +class QueueLimits: + """Bounds for pending samples and complete plans awaiting execution.""" + + lookahead_samples: int + prepared_batches: int = 2 + + def __post_init__(self) -> None: + _positive_int(self.lookahead_samples, "lookahead_samples", (1 << 32) - 5) + _positive_int(self.prepared_batches, "prepared_batches", 1024) + + +@dataclass(frozen=True, slots=True) +class MetadataCache: + """Cache capacities used by a metadata provider.""" + + array_entries: int = 256 + shard_index_entries: int = 8192 + + def __post_init__(self) -> None: + _positive_int(self.array_entries, "array_entries") + _positive_int(self.shard_index_entries, "shard_index_entries") + + +@dataclass(frozen=True, slots=True) +class PlanLimits: + """Per-batch planning limits; max_chunks counts uses before deduplication.""" + + max_chunks: int = 16384 + max_chunk_bytes: int = 2 << 20 + max_shards_per_sample: int = 64 + max_plan_bytes: int = 64 << 20 + + def __post_init__(self) -> None: + _positive_int(self.max_chunks, "max_chunks", 16384) + _positive_int(self.max_chunk_bytes, "max_chunk_bytes") + _positive_int(self.max_shards_per_sample, "max_shards_per_sample") + _positive_int(self.max_plan_bytes, "max_plan_bytes", (1 << 64) - 1) + + +@dataclass(frozen=True, slots=True) +class CpuLimits: + """Bound input, decoded data, codec workspace, and output storage.""" + + max_memory_bytes: int + decode_workers: int = 8 + max_encoded_chunk_bytes: int = 4 << 20 + max_decoded_chunk_bytes: int = 2 << 20 + + def __post_init__(self) -> None: + _positive_int(self.max_memory_bytes, "max_memory_bytes", (1 << 64) - 1) + _positive_int(self.decode_workers, "decode_workers", _native.MAX_IO_THREADS) + _positive_int(self.max_encoded_chunk_bytes, "max_encoded_chunk_bytes") + _positive_int(self.max_decoded_chunk_bytes, "max_decoded_chunk_bytes") + + +@dataclass(frozen=True, slots=True) +class CudaLimits: + """GPU buffer limits, wave geometry, and codec-layout cache capacity.""" + + max_gpu_memory_bytes: int + chunk_layout_entries: int = 256 + max_chunk_bytes: int = _native.DEFAULT_CHUNK_UNCOMPRESSED_BYTES + max_read_bytes: int = _native.DEFAULT_READ_OP_MAX_BYTES + host_buffer_waves: int = _native.DEFAULT_HOST_BUFFER_WAVES + max_chunks_per_wave: int = _native.DEFAULT_MAX_CHUNKS_PER_WAVE + max_substreams_per_chunk: int = _native.DEFAULT_MAX_SUBSTREAMS_PER_CHUNK + + def __post_init__(self) -> None: + _positive_int(self.chunk_layout_entries, "chunk_layout_entries") + _positive_int(self.max_gpu_memory_bytes, "max_gpu_memory_bytes", (1 << 64) - 1) + _positive_int(self.max_chunk_bytes, "max_chunk_bytes") + _positive_int(self.max_read_bytes, "max_read_bytes") + _positive_int( + self.host_buffer_waves, "host_buffer_waves", _native.MAX_HOST_BUFFER_WAVES + ) + if self.host_buffer_waves < _native.N_WAVES: + raise ValueError(f"host_buffer_waves must be at least {_native.N_WAVES}") + _positive_int( + self.max_chunks_per_wave, + "max_chunks_per_wave", + _native.HARD_MAX_CHUNKS_PER_WAVE, + ) + _positive_int( + self.max_substreams_per_chunk, + "max_substreams_per_chunk", + _native.HARD_MAX_SUBSTREAMS_PER_CHUNK, + ) + + +class FileReader: + """Read bulk encoded chunk data through a bounded filesystem I/O queue.""" + + __slots__ = ("_native",) + + def __init__(self, *, workers: int = 8, max_inflight_reads: int = 4096) -> None: + self._native = _component( + _native.create_reader, + _positive_int(workers, "workers", _native.MAX_IO_THREADS), + _positive_int(max_inflight_reads, "max_inflight_reads"), + ) + + +class FileMetadataReader: + """Configure a separate asynchronous queue for filesystem metadata reads.""" + + __slots__ = ("_native",) + + def __init__( + self, *, concurrency: int = 64, latency: LatencyModel | None = None + ) -> None: + latency = latency or LatencyModel() + if not isinstance(latency, LatencyModel): + raise TypeError("latency must be a LatencyModel") + self._native = _component( + _native.create_metadata_reader, + _positive_int(concurrency, "concurrency", 4096), + latency.baseline_ns, + latency.lognormal_mu_ln_ns, + latency.lognormal_sigma_ln_ns, + latency.cap_ns, + latency.seed, + ) + + +class ZarrMetadata: + """Provide Zarr v3 array metadata and shard indexes for sample URIs.""" + + __slots__ = ("_native", "cache", "reader") + + def __init__(self, *, reader: FileMetadataReader, cache: MetadataCache) -> None: + self.reader = reader + self.cache = cache + self._native = _component( + _native.create_metadata, + reader._native, + cache.array_entries, + cache.shard_index_entries, + ) + + +class ChunkPlanner: + """Prepare owned chunk plans from rectangular queries and Zarr metadata.""" + + __slots__ = ("_native", "limits", "metadata") + + def __init__( + self, *, metadata: ZarrMetadata, limits: PlanLimits | None = None + ) -> None: + limits = limits if limits is not None else PlanLimits() + self.metadata = metadata + self.limits = limits + self._native = _component( + _native.create_planner, + metadata._native, + limits.max_chunks, + limits.max_chunk_bytes, + limits.max_shards_per_sample, + limits.max_plan_bytes, + ) + + +class CpuExecutor: + """Decode and assemble batches in ordinary RAM without using CUDA.""" + + __slots__ = ("_native", "limits", "reader") + + def __init__(self, *, reader: FileReader, limits: CpuLimits) -> None: + self.reader = reader + self.limits = limits + self._native = _component( + _native.create_cpu_executor, + reader._native, + limits.decode_workers, + limits.max_encoded_chunk_bytes, + limits.max_decoded_chunk_bytes, + limits.max_memory_bytes, + ) + + +class CudaExecutor: + """Decode and assemble batches on a CUDA device.""" + + __slots__ = ("_native", "device", "limits", "reader") + + def __init__( + self, + *, + reader: FileReader, + limits: CudaLimits, + device: int | None = None, + numa_strategy: NumaStrategy | str | int = NumaStrategy.AUTO, + numa_node: int = -1, + enable_gds: bool | None = None, + ) -> None: + self.reader = reader + self.limits = limits + self.device = device + self._native = _component( + _native.create_cuda_executor, + reader._native, + -1 if device is None else operator.index(device), + limits.max_gpu_memory_bytes, + limits.max_chunk_bytes, + limits.max_read_bytes, + limits.max_chunks_per_wave, + limits.max_substreams_per_chunk, + limits.host_buffer_waves, + limits.chunk_layout_entries, + int(NumaStrategy.coerce(numa_strategy)), + numa_node, + _gds_to_native(enable_gds), + ) + + @dataclass(frozen=True, slots=True) class BatchInfo: - """Snapshot of the on-device batch geometry.""" + """Batch geometry and memory location. CPU data is ready for host access.""" device_ptr: int shape: tuple[int, ...] dtype: Dtype ready_stream: int batch_id: int + device_type: DeviceType = DeviceType.CUDA + device_id: int = 0 + + @property + def data(self) -> int: + return self.device_ptr @classmethod def _from_native(cls, info: dict[str, Any]) -> BatchInfo: @@ -740,6 +1016,8 @@ def _from_native(cls, info: dict[str, Any]) -> BatchInfo: dtype=Dtype.coerce(info["dtype"]), ready_stream=info["ready_stream"], batch_id=info["batch_id"], + device_type=DeviceType(info["device_type"]), + device_id=info["device_id"], ) @@ -804,6 +1082,7 @@ class Stats: reads_issued: int worker_steps: int gpu_bytes_committed: int + host_bytes_committed: int = 0 @classmethod def _from_native(cls, st: dict[str, Any]) -> Stats: @@ -842,6 +1121,7 @@ def _from_native(cls, st: dict[str, Any]) -> Stats: reads_issued=st["reads_issued"], worker_steps=st["worker_steps"], gpu_bytes_committed=st["gpu_bytes_committed"], + host_bytes_committed=st["host_bytes_committed"], ) @@ -919,29 +1199,23 @@ def _coerce_cuda_event_handle(event: object) -> int | None: class Batch: - """A batch of samples on the device, ready for consumption. + """A contiguous batch in CPU or CUDA memory, ready for consumption. - Use as a context manager to release the slot back to the pool:: + Use a context manager to release the batch reference:: - with d.pop() as batch: - x = torch.from_dlpack(batch) + with pipeline.pop() as batch: + tensor = torch.from_dlpack(batch) - The DLPack capsule (``batch.__dlpack__()``) keeps the underlying - storage alive as long as the consumer holds it; releasing the - Batch object while a tensor still views it is safe. + A DLPack consumer retains the storage independently. Its view stays valid + after batch release and pipeline close; the pool cannot reuse that buffer + until all consumers release it. - **Deferred release.** If the consumer kicks off an async D2D copy on - a side stream, the default ``with`` block forces a host-side - ``cuStreamSynchronize`` on the producer stream before the slot is - reused. To avoid that block, call :meth:`release` explicitly with - the consumer's stream or event — damacy will stream-wait on it - before re-assembling into the slot's buffer:: + For asynchronous CUDA work on a side stream, release with the consumer's + stream or event so subsequent output writes wait for that work:: - batch = d.pop() - tensor = torch.empty_like(...) # on side_stream - with torch.cuda.stream(side_stream): - tensor.copy_(torch.from_dlpack(batch)) - batch.release(event=side_stream) # no host sync + batch.release(event=side_stream) + + CPU batches accept only immediate release (``event=None``). """ __slots__ = ("_native",) @@ -951,7 +1225,7 @@ def __init__(self, native_batch: _native.Batch) -> None: @property def info(self) -> BatchInfo: - """Snapshot of the on-device batch geometry. Raises after release.""" + """Snapshot of batch geometry and device. Raises after release.""" return BatchInfo._from_native(self._native.info) def release( @@ -959,12 +1233,12 @@ def release( *, event: object | None = None, ) -> None: - """Return the slot to the pool. Idempotent. + """Release this handle. Exported tensors keep their buffers. Idempotent. Args: - event: If ``None`` (default), the slot is freed immediately; - damacy may reuse the buffer right away, so callers must - have host-synced any work that reads it. Otherwise the + event: If ``None`` (default), storage becomes reusable when the + last exported tensor and batch handle are released. Callers + must finish work reading the buffer before then. Otherwise the slot reuse waits on the supplied CUDA event before damacy's assemble kernel writes the buffer again — the host returns at once. Accepted forms: @@ -1103,28 +1377,24 @@ def _warn_if_multi_gpu_implicit(cfg_device: int | None, bound: int) -> None: class Pipeline: - """Streaming GPU data pipeline. Drive :meth:`push`, :meth:`pop`. - Stages are plan → host I/O → H2D copy → on-device - decompress → assemble; output batches are double-buffered (B=2) - and waves are double-buffered internally. - - A CUcontext must be current on the calling thread when this is - constructed; PyTorch sets one up implicitly. For bare-Python use, - call :func:`damacy._native.cuda_init_primary` once first. + """Load batches using an injected planner and CPU or CUDA executor. - Constructed from a :class:`Config`:: + ``planner`` resolves source metadata into owned chunk plans. ``executor`` + reads, decodes, and assembles them. ``output`` defines the batch tensor; + ``queues`` bounds preparation. Components serve one active pipeline and + may be reused after it closes. Exported tensors retain their storage. - cfg = damacy.Config(samples_per_batch=8, ...) - with damacy.Pipeline(cfg) as p: - ... - - Resource caps are fixed at construction; nothing grows after that. + ``Pipeline(Config(...))`` composes the CUDA pipeline for existing callers. + For CUDA, pass an explicit executor device or make a CUDA context current + before constructing the pipeline. CPU execution requires neither. """ __slots__ = ( "_closed", + "_components", "_config", "_native", + "_output", "_pending", "_pending_buf", "_pop_done", @@ -1132,43 +1402,84 @@ class Pipeline: "_pop_lock", "_pop_result", "_pop_thread", + "_pop_timeout_s", + "_queues", ) - def __init__(self, config: Config) -> None: - try: - self._native = _native.Pipeline( - samples_per_batch=config.samples_per_batch, - lookahead_samples=config.lookahead_samples, - dtype=int(config.dtype), # already coerced by Config.__init__ - max_chunk_uncompressed_bytes=config.max_chunk_uncompressed_bytes, - max_read_op_bytes=config.max_read_op_bytes, - max_gpu_memory_bytes=config.max_gpu_memory_bytes, - host_buffer_waves=config.host_buffer_waves, - max_chunks_per_wave=config.max_chunks_per_wave, - max_substreams_per_chunk=config.max_substreams_per_chunk, - n_io_threads=config.n_io_threads, - metadata_io_concurrency=config.metadata_io_concurrency, - n_array_meta_cache=config.n_array_meta_cache, - n_shard_index_cache=config.n_shard_index_cache, - n_chunk_layout_cache=config.n_chunk_layout_cache, - max_shards_per_sample=config.max_shards_per_sample, - sample_shape=tuple(config.sample_shape), - device=-1 if config.device is None else int(config.device), - enable_gds=_gds_to_native(config.enable_gds), - numa_strategy=int(config.numa_strategy), - numa_node=config.numa_node, - metadata_latency_baseline_ns=config.metadata_latency.baseline_ns, - metadata_latency_lognormal_mu_ln_ns=( - config.metadata_latency.lognormal_mu_ln_ns - ), - metadata_latency_lognormal_sigma_ln_ns=( - config.metadata_latency.lognormal_sigma_ln_ns - ), - metadata_latency_cap_ns=config.metadata_latency.cap_ns, - metadata_latency_seed=config.metadata_latency.seed, + def __init__( + self, + config: Config | None = None, + *, + planner: ChunkPlanner | None = None, + executor: CpuExecutor | CudaExecutor | None = None, + output: BatchSpec | None = None, + queues: QueueLimits | None = None, + pop_timeout_s: float | None = 30.0, + ) -> None: + if config is not None: + if any(x is not None for x in (planner, executor, output, queues)): + raise TypeError("supply either Config or pipeline components") + output = BatchSpec( + config.samples_per_batch, config.sample_shape, config.dtype ) + queues = QueueLimits(config.lookahead_samples) + pop_timeout_s = config.pop_timeout_s + elif planner is None or executor is None or output is None or queues is None: + raise TypeError("planner, executor, output, and queues are required") + if pop_timeout_s is not None and ( + not math.isfinite(pop_timeout_s) or pop_timeout_s <= 0 + ): + raise ValueError("pop_timeout_s must be positive and finite, or None") + try: + if config is not None: + self._native = _native.Pipeline( + samples_per_batch=config.samples_per_batch, + lookahead_samples=config.lookahead_samples, + dtype=int(config.dtype), # already coerced by Config.__init__ + max_chunk_uncompressed_bytes=config.max_chunk_uncompressed_bytes, + max_read_op_bytes=config.max_read_op_bytes, + max_gpu_memory_bytes=config.max_gpu_memory_bytes, + host_buffer_waves=config.host_buffer_waves, + max_chunks_per_wave=config.max_chunks_per_wave, + max_substreams_per_chunk=config.max_substreams_per_chunk, + n_io_threads=config.n_io_threads, + metadata_io_concurrency=config.metadata_io_concurrency, + n_array_meta_cache=config.n_array_meta_cache, + n_shard_index_cache=config.n_shard_index_cache, + n_chunk_layout_cache=config.n_chunk_layout_cache, + max_shards_per_sample=config.max_shards_per_sample, + sample_shape=tuple(config.sample_shape), + device=-1 if config.device is None else int(config.device), + enable_gds=_gds_to_native(config.enable_gds), + numa_strategy=int(config.numa_strategy), + numa_node=config.numa_node, + metadata_latency_baseline_ns=config.metadata_latency.baseline_ns, + metadata_latency_lognormal_mu_ln_ns=( + config.metadata_latency.lognormal_mu_ln_ns + ), + metadata_latency_lognormal_sigma_ln_ns=( + config.metadata_latency.lognormal_sigma_ln_ns + ), + metadata_latency_cap_ns=config.metadata_latency.cap_ns, + metadata_latency_seed=config.metadata_latency.seed, + ) + else: + assert planner is not None and executor is not None + self._native = _native.compose_pipeline( + planner._native, + executor._native, + output.shape, + output.samples, + int(output.dtype), + queues.lookahead_samples, + queues.prepared_batches, + ) except _native.DamacyError as exc: _reraise_typed(exc) + self._output = output + self._queues = queues + self._pop_timeout_s = pop_timeout_s + self._components = (planner, executor) self._closed = False self._config = config # User-side queue of pending sample iterators. push() appends @@ -1188,20 +1499,36 @@ def __init__(self, config: Config) -> None: self._pop_result: _native.Batch | None = None self._pop_err: BaseException | None = None bound = self._native.device - if not _warn_if_local_rank_disagrees(config.device, bound): - _warn_if_multi_gpu_implicit(config.device, bound) + if bound >= 0: + configured_device = ( + config.device + if config is not None + else executor.device + if isinstance(executor, CudaExecutor) + else None + ) + if not _warn_if_local_rank_disagrees(configured_device, bound): + _warn_if_multi_gpu_implicit(configured_device, bound) @property def device(self) -> int: - """CUDA device index this pipeline is bound to.""" + """CUDA device index, or -1 for CPU execution.""" self._check_open() return int(self._native.device) @property - def config(self) -> Config: + def config(self) -> Config | None: """The :class:`Config` this loader was built from.""" return self._config + @property + def output(self) -> BatchSpec: + return self._output + + @property + def queues(self) -> QueueLimits: + return self._queues + # ---- lifecycle --------------------------------------------------- def _check_open(self) -> None: @@ -1218,11 +1545,17 @@ def close(self) -> None: on the pipeline raise :class:`ShutdownError`.""" if not self._closed: self._closed = True - del self._native + self._native.shutdown() t = self._pop_thread if t is not None: - t.join() # damacy_destroy already woke it with SHUTDOWN + t.join() self._pop_thread = None + self._pop_result = None + self._pop_err = None + self._pending.clear() + self._pending_buf.clear() + del self._native + self._components = (None, None) def __enter__(self) -> Self: return self @@ -1243,7 +1576,7 @@ def push(self, samples: Iterable[Sample]) -> None: generator, infinite generator, …); large or unbounded sources are pulled lazily as :meth:`pop` frees space. - Local validation (shape/rank against ``Config.sample_shape``) + Local validation (shape/rank against ``Pipeline.output.shape``) raises the matching :class:`DamacyError` subclass here and discards the offending iterator. Errors that depend on store contents — :class:`NotFound`, :class:`DtypeMismatch`, @@ -1255,7 +1588,7 @@ def push(self, samples: Iterable[Sample]) -> None: :class:`ShutdownError`. Batching is ``drop_last=True``: only complete batches of - ``Config.samples_per_batch`` are emitted, so trailing samples + ``Pipeline.output.samples`` are emitted, so trailing samples beyond the last whole multiple are never returned. (Emitting the ragged final batch is not yet supported — issue #139.) """ @@ -1272,7 +1605,7 @@ def _drain_pending(self) -> None: a single head iterator), not re-wrapped onto ``self._pending[0]`` — successive backpressure events leave the buffer flat instead of nesting ``itertools.chain`` layers.""" - cap = self._config.lookahead_samples + cap = self._queues.lookahead_samples # Top up buffer from the head iterator. Buffer is only ever # filled from one iterator at a time, so on push failure we know # exactly which iterator to drop. Each refill pulls at most one @@ -1322,7 +1655,7 @@ def _drain_pending(self) -> None: self._pending.popleft() def pop(self) -> Batch: - """Block until the next batch is on-device-ready. Returns a + """Block until the next batch is ready. Returns a :class:`Batch` you can hand to ``torch.from_dlpack`` (or any DLPack consumer) — preferably inside a ``with`` block. @@ -1345,7 +1678,7 @@ def pop(self) -> Batch: # not the secondary SHUTDOWN raised by re-pushing into a terminal. with contextlib.suppress(ShutdownError): self._drain_pending() - timeout = self._config.pop_timeout_s + timeout = self._pop_timeout_s if timeout is None: try: return Batch(self._native.pop()) @@ -1403,7 +1736,7 @@ def pending(self) -> bool: def batches(self, n: int) -> Iterator[Batch]: """Pop *n* batches as an iterator. Each call to :meth:`pop` - blocks until that batch is on-device-ready. + blocks until that batch is ready. Pair with a ``with`` block so the slot is released:: diff --git a/python/damacy/_api.c b/python/damacy/_api.c index 5e5d2378..621f6706 100644 --- a/python/damacy/_api.c +++ b/python/damacy/_api.c @@ -17,7 +17,9 @@ #include "_api.h" #include "damacy.h" +#ifdef DAMACY_HAS_CUDA #include +#endif #include #include #include @@ -216,7 +218,8 @@ raise_status(enum damacy_status s, const char* what) typedef struct { - PyObject_HEAD struct damacy* handle; // strong ref while not destroyed + PyObject_HEAD struct damacy* handle; + PyObject* dependencies; } PipelineObj; typedef struct @@ -315,9 +318,15 @@ Batch_info(BatchObj* self, void* Py_UNUSED(closure)) PyTuple_SET_ITEM(shape, i, v); } - return Py_BuildValue("{s:K,s:N,s:s,s:K,s:K}", + return Py_BuildValue("{s:K,s:K,s:i,s:i,s:N,s:s,s:K,s:K}", + "data", + (unsigned long long)(uintptr_t)info.data, "device_ptr", - (unsigned long long)(uintptr_t)info.device_ptr, + (unsigned long long)(uintptr_t)info.data, + "device_type", + (int)info.device_type, + "device_id", + info.device_id, "shape", shape, "dtype", @@ -356,21 +365,16 @@ struct dlpack_payload DLManagedTensor mt_v0; DLManagedTensorVersioned mt_v1; int64_t shape[DAMACY_MAX_RANK + 1]; // referenced by dl_tensor.shape - PyObject* batch; // strong ref; dropped by deleter + struct damacy_batch* handle; }; -// Drop the batch ref and free the payload under the GIL. Shared by both -// v0 and v1 deleters. PyMem_Free uses pymalloc which itself requires the -// GIL, so the free stays inside the GIL window. static void dlpack_payload_free(struct dlpack_payload* p) { if (!p) return; - PyGILState_STATE g = PyGILState_Ensure(); - Py_XDECREF(p->batch); - PyMem_Free(p); - PyGILState_Release(g); + damacy_batch_release(p->handle); + free(p); } static void @@ -409,6 +413,7 @@ dlpack_capsule_destructor(PyObject* capsule) } } +#ifdef DAMACY_HAS_CUDA static int sync_streams_for_consumer(void* producer_stream_v, PyObject* stream_obj) { @@ -481,6 +486,8 @@ sync_streams_for_consumer(void* producer_stream_v, PyObject* stream_obj) return 0; } +#endif + // Parse max_version per the array-API DLPack protocol. Returns 0 on // success, -1 on error (PyErr set). *out_major / *out_minor land at the // requested version, or (0,0) when the consumer didn't specify one. @@ -539,7 +546,6 @@ Batch_dlpack(BatchObj* self, PyObject* args, PyObject* kw) if (!PyArg_ParseTupleAndKeywords( args, kw, "|OOOO", kws, &stream_obj, &max_version, &dl_device, ©)) return NULL; - (void)dl_device; if (copy != Py_None && PyObject_IsTrue(copy)) { PyErr_SetString(PyExc_BufferError, "damacy DLPack: copy=True not supported"); @@ -568,30 +574,43 @@ Batch_dlpack(BatchObj* self, PyObject* args, PyObject* kw) } (void)bpe; - // Resolve device id from the device pointer. - int dev_id = 0; - unsigned int ord = 0; - if (cuPointerGetAttribute(&ord, - CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL, - (CUdeviceptr)info.device_ptr) == CUDA_SUCCESS) - dev_id = (int)ord; - - if (sync_streams_for_consumer(info.ready_stream, stream_obj) != 0) + if (dl_device != Py_None) { + int requested_type, requested_id; + if (!PyArg_ParseTuple(dl_device, "ii", &requested_type, &requested_id)) + return NULL; + if (requested_type != (int)info.device_type || + requested_id != info.device_id) { + PyErr_SetString(PyExc_BufferError, + "requested device differs from the batch device"); + return NULL; + } + } + if (info.device_type == DAMACY_DEVICE_CPU) { + if (stream_obj != Py_None) { + PyErr_SetString(PyExc_ValueError, "CPU tensors require stream=None"); + return NULL; + } + } +#ifdef DAMACY_HAS_CUDA + else if (info.ready_stream && + sync_streams_for_consumer(info.ready_stream, stream_obj) != 0) return NULL; +#endif - struct dlpack_payload* p = PyMem_Calloc(1, sizeof *p); + struct dlpack_payload* p = calloc(1, sizeof *p); if (!p) return PyErr_NoMemory(); for (int i = 0; i < info.rank; ++i) p->shape[i] = info.shape[i]; - Py_INCREF(self); - p->batch = (PyObject*)self; + p->handle = self->handle; + damacy_batch_retain(p->handle); DLTensor dl = { - .data = info.device_ptr, - .device = (DLDevice){ .device_type = kDLCUDA, .device_id = dev_id }, + .data = info.data, + .device = (DLDevice){ .device_type = info.device_type, + .device_id = info.device_id }, .ndim = (int32_t)info.rank, .dtype = dlt, .shape = p->shape, @@ -615,8 +634,7 @@ Batch_dlpack(BatchObj* self, PyObject* args, PyObject* kw) cap = PyCapsule_New(&p->mt_v0, "dltensor", dlpack_capsule_destructor); } if (!cap) { - Py_DECREF(self); - PyMem_Free(p); + dlpack_payload_free(p); return NULL; } return cap; @@ -628,22 +646,16 @@ Batch_dlpack_device(BatchObj* self, PyObject* Py_UNUSED(ignored)) RETURN_IF_DESTROYED(self, "Batch has been released"); struct damacy_batch_info info; damacy_batch_info(self->handle, &info); - unsigned int ord = 0; - if (cuPointerGetAttribute(&ord, - CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL, - (CUdeviceptr)info.device_ptr) != CUDA_SUCCESS) - ord = 0; - return Py_BuildValue("(ii)", (int)kDLCUDA, (int)ord); + return Py_BuildValue("(ii)", (int)info.device_type, info.device_id); } static PyMethodDef Batch_methods[] = { { "release", (PyCFunction)(void (*)(void))Batch_release, METH_VARARGS | METH_KEYWORDS, - "release(event=None): return the slot to the pool. With event=None " - "(default), release is immediate. With event set to an integer CUevent " - "handle, damacy stream-waits on it before reusing the slot's buffer — " - "the host returns immediately. Idempotent." }, + "release(event=None): release this batch reference. DLPack consumers " + "keep the buffer alive. A CUDA event delays reuse until that event " + "completes. Idempotent." }, { "__dlpack__", (PyCFunction)(void (*)(void))Batch_dlpack, METH_VARARGS | METH_KEYWORDS, @@ -654,7 +666,7 @@ static PyMethodDef Batch_methods[] = { { "__dlpack_device__", (PyCFunction)Batch_dlpack_device, METH_NOARGS, - "DLPack device tuple: (kDLCUDA=2, ordinal)." }, + "DLPack device tuple: (1, 0) for CPU or (2, ordinal) for CUDA." }, { NULL, NULL, 0, NULL }, }; @@ -695,6 +707,10 @@ batch_new(PipelineObj* parent, struct damacy_batch* handle) static int Pipeline_init(PipelineObj* self, PyObject* args, PyObject* kw) { + if (self->handle) { + PyErr_SetString(PyExc_RuntimeError, "Pipeline is already initialized"); + return -1; + } // kws[] / format string / variable list mirror struct damacy_config — // keep all three in sync when adding a field. static char* kws[] = { "samples_per_batch", @@ -739,7 +755,8 @@ Pipeline_init(PipelineObj* self, PyObject* args, PyObject* kw) PyObject* sample_shape_obj = NULL; unsigned int host_buffer_waves = DAMACY_DEFAULT_HOST_BUFFER_WAVES; unsigned int max_chunks_per_wave = DAMACY_DEFAULT_MAX_CHUNKS_PER_WAVE; - unsigned int max_substreams_per_chunk = DAMACY_DEFAULT_MAX_SUBSTREAMS_PER_CHUNK; + unsigned int max_substreams_per_chunk = + DAMACY_DEFAULT_MAX_SUBSTREAMS_PER_CHUNK; unsigned long long max_read_op_bytes = td.max_read_op_bytes; int device = -1; int enable_gds = DAMACY_GDS_AUTO; @@ -896,6 +913,7 @@ Pipeline_dealloc(PipelineObj* self) self->handle = NULL; WITH_GIL_RELEASED(damacy_destroy(d)); } + Py_XDECREF(self->dependencies); Py_TYPE(self)->tp_free((PyObject*)self); } @@ -996,7 +1014,11 @@ Pipeline_pop(PipelineObj* self, PyObject* Py_UNUSED(ignored)) WITH_GIL_RELEASED(s = damacy_pop(self->handle, &b)); if (s != DAMACY_OK) return raise_status(s, "pop"); - return (PyObject*)batch_new(self, b); + BatchObj* batch = batch_new(self, b); + if (!batch) { + WITH_GIL_RELEASED(damacy_batch_release(b)); + } + return (PyObject*)batch; } static PyObject* @@ -1088,6 +1110,7 @@ Pipeline_stats(PipelineObj* self, PyObject* Py_UNUSED(ignored)) { "chunks_dispatched", st.chunks_dispatched }, { "reads_issued", st.reads_issued }, { "gpu_bytes_committed", st.gpu_bytes_committed }, + { "host_bytes_committed", st.host_bytes_committed }, }; for (size_t i = 0; i < sizeof counters / sizeof counters[0]; ++i) if (dict_set_steal(d, @@ -1125,7 +1148,19 @@ static PyGetSetDef Pipeline_getset[] = { { NULL, NULL, NULL, NULL, NULL }, }; +static PyObject* +Pipeline_shutdown(PipelineObj* self, PyObject* Py_UNUSED(ignored)) +{ + RETURN_IF_DESTROYED(self, "Pipeline has been destroyed"); + WITH_GIL_RELEASED(damacy_shutdown(self->handle)); + Py_RETURN_NONE; +} + static PyMethodDef Pipeline_methods[] = { + { "shutdown", + (PyCFunction)Pipeline_shutdown, + METH_NOARGS, + "Stop work and retain exported buffers." }, { "push", (PyCFunction)Pipeline_push, METH_O, @@ -1160,6 +1195,35 @@ PyTypeObject PipelineType = { .tp_getset = Pipeline_getset, }; +PyObject* +api_raise_status(enum damacy_status status, const char* what) +{ + return raise_status(status, what); +} + +PyObject* +api_pipeline_from_components(struct damacy_planner* planner, + struct damacy_executor* executor, + const struct damacy_batch_spec* output, + const struct damacy_queue_limits* queues, + PyObject* dependencies) +{ + PipelineObj* self = (PipelineObj*)PipelineType.tp_alloc(&PipelineType, 0); + if (!self) + return NULL; + Py_INCREF(dependencies); + self->dependencies = dependencies; + enum damacy_status status; + Py_BEGIN_ALLOW_THREADS status = + damacy_pipeline_create(planner, executor, output, queues, &self->handle); + Py_END_ALLOW_THREADS if (status != DAMACY_OK) + { + Py_DECREF(self); + return raise_status(status, "create"); + } + return (PyObject*)self; +} + // ---------- registration ---------- int diff --git a/python/damacy/_api.h b/python/damacy/_api.h index 27d72e58..1478e108 100644 --- a/python/damacy/_api.h +++ b/python/damacy/_api.h @@ -1,18 +1,17 @@ #pragma once -// Internal: registers the Pipeline and Batch types on a module object. -// Returns 0 on success, -1 with a Python error set on failure. - -#define PY_SSIZE_T_CLEAN +#include "damacy_pipeline.h" #include -#ifdef __cplusplus -extern "C" -{ -#endif - - int api_register_types(PyObject* module); - -#ifdef __cplusplus -} -#endif +int +api_register_types(PyObject* module); +int +components_register(PyObject* module); +PyObject* +api_raise_status(enum damacy_status status, const char* what); +PyObject* +api_pipeline_from_components(struct damacy_planner* planner, + struct damacy_executor* executor, + const struct damacy_batch_spec* output, + const struct damacy_queue_limits* queues, + PyObject* dependencies); diff --git a/python/damacy/_components.c b/python/damacy/_components.c new file mode 100644 index 00000000..16677b4d --- /dev/null +++ b/python/damacy/_components.c @@ -0,0 +1,320 @@ +#define PY_SSIZE_T_CLEAN +#include + +#include "_api.h" +#include "damacy_pipeline.h" + +enum component_kind +{ + READER, + METADATA_READER, + METADATA, + PLANNER, + EXECUTOR +}; + +struct component +{ + enum component_kind kind; + void* value; + PyObject* dependencies; +}; + +static const char* names[] = { "damacy.Reader", + "damacy.MetadataReader", + "damacy.Metadata", + "damacy.Planner", + "damacy.Executor" }; + +static void +destroy_value(enum component_kind kind, void* value) +{ + switch (kind) { + case READER: + damacy_reader_destroy(value); + break; + case METADATA_READER: + damacy_metadata_reader_destroy(value); + break; + case METADATA: + damacy_metadata_destroy(value); + break; + case PLANNER: + damacy_planner_destroy(value); + break; + case EXECUTOR: + damacy_executor_destroy(value); + break; + } +} + +static void +component_destroy(PyObject* capsule) +{ + struct component* component = + PyCapsule_GetPointer(capsule, PyCapsule_GetName(capsule)); + if (!component) + return; + Py_BEGIN_ALLOW_THREADS destroy_value(component->kind, component->value); + Py_END_ALLOW_THREADS Py_XDECREF(component->dependencies); + PyMem_Free(component); +} + +static PyObject* +component_new(enum component_kind kind, void* value, PyObject* dependencies) +{ + struct component* component = PyMem_Calloc(1, sizeof(*component)); + if (!component) { + destroy_value(kind, value); + return PyErr_NoMemory(); + } + component->kind = kind; + component->value = value; + Py_XINCREF(dependencies); + component->dependencies = dependencies; + PyObject* capsule = PyCapsule_New(component, names[kind], component_destroy); + if (!capsule) { + Py_XDECREF(dependencies); + destroy_value(kind, value); + PyMem_Free(component); + } + return capsule; +} + +static void* +component_value(PyObject* capsule, enum component_kind kind) +{ + struct component* component = PyCapsule_GetPointer(capsule, names[kind]); + return component ? component->value : NULL; +} + +static PyObject* +create_reader(PyObject* self, PyObject* args) +{ + (void)self; + unsigned int workers, reads; + if (!PyArg_ParseTuple(args, "II", &workers, &reads)) + return NULL; + struct damacy_reader* reader = NULL; + enum damacy_status status; + Py_BEGIN_ALLOW_THREADS status = + damacy_file_reader_create(workers, reads, &reader); + Py_END_ALLOW_THREADS if (status != DAMACY_OK) return api_raise_status( + status, "create reader"); + return component_new(READER, reader, NULL); +} + +static PyObject* +create_metadata_reader(PyObject* self, PyObject* args) +{ + (void)self; + unsigned int concurrency; + unsigned long long baseline, cap, seed; + double mu, sigma; + if (!PyArg_ParseTuple( + args, "IKddKK", &concurrency, &baseline, &mu, &sigma, &cap, &seed)) + return NULL; + struct damacy_metadata_reader* reader = NULL; + struct damacy_latency_model latency = { .baseline_ns = baseline, + .lognormal_mu_ln_ns = mu, + .lognormal_sigma_ln_ns = sigma, + .cap_ns = cap, + .seed = seed }; + enum damacy_status status = + damacy_file_metadata_reader_create(concurrency, &latency, &reader); + if (status != DAMACY_OK) + return api_raise_status(status, "create metadata reader"); + return component_new(METADATA_READER, reader, NULL); +} + +static PyObject* +create_metadata(PyObject* self, PyObject* args) +{ + (void)self; + PyObject* reader_object; + struct damacy_metadata_cache_config cache; + if (!PyArg_ParseTuple(args, + "OII", + &reader_object, + &cache.array_entries, + &cache.shard_entries)) + return NULL; + struct damacy_metadata_reader* reader = + component_value(reader_object, METADATA_READER); + if (!reader) + return NULL; + struct damacy_metadata* metadata = NULL; + enum damacy_status status = + damacy_zarr_metadata_create(reader, &cache, &metadata); + if (status != DAMACY_OK) + return api_raise_status(status, "create metadata"); + return component_new(METADATA, metadata, reader_object); +} + +static PyObject* +create_planner(PyObject* self, PyObject* args) +{ + (void)self; + PyObject* metadata_object; + unsigned long long bytes; + struct damacy_plan_limits limits; + if (!PyArg_ParseTuple(args, + "OIIIK", + &metadata_object, + &limits.max_chunks, + &limits.max_chunk_bytes, + &limits.max_shards_per_sample, + &bytes)) + return NULL; + limits.max_plan_bytes = bytes; + struct damacy_metadata* metadata = component_value(metadata_object, METADATA); + if (!metadata) + return NULL; + struct damacy_planner* planner = NULL; + enum damacy_status status = + damacy_chunk_planner_create(metadata, &limits, &planner); + if (status != DAMACY_OK) + return api_raise_status(status, "create planner"); + return component_new(PLANNER, planner, metadata_object); +} + +static PyObject* +create_cpu_executor(PyObject* self, PyObject* args) +{ + (void)self; + PyObject* reader_object; + unsigned long long bytes; + struct damacy_cpu_config config; + if (!PyArg_ParseTuple(args, + "OIIIK", + &reader_object, + &config.decode_workers, + &config.max_encoded_chunk_bytes, + &config.max_decoded_chunk_bytes, + &bytes)) + return NULL; + config.max_memory_bytes = bytes; + struct damacy_reader* reader = component_value(reader_object, READER); + if (!reader) + return NULL; + struct damacy_executor* executor = NULL; + enum damacy_status status = + damacy_cpu_executor_create(reader, &config, &executor); + if (status != DAMACY_OK) + return api_raise_status(status, "create CPU executor"); + return component_new(EXECUTOR, executor, reader_object); +} + +static PyObject* +create_cuda_executor(PyObject* self, PyObject* args) +{ + (void)self; + PyObject* reader_object; + unsigned long long memory, read_bytes; + unsigned int buffers; + int numa, gds; + struct damacy_cuda_config config = { 0 }; + if (!PyArg_ParseTuple(args, + "OiKIKIIIIiii", + &reader_object, + &config.device, + &memory, + &config.max_chunk_bytes, + &read_bytes, + &config.max_chunks_per_wave, + &config.max_substreams_per_chunk, + &buffers, + &config.chunk_layout_entries, + &numa, + &config.numa_node, + &gds)) + return NULL; + if (buffers > UINT8_MAX) { + PyErr_SetString(PyExc_ValueError, "host_buffer_waves is too large"); + return NULL; + } + config.max_gpu_memory_bytes = memory; + config.max_read_bytes = read_bytes; + config.host_buffer_waves = (uint8_t)buffers; + config.numa_strategy = (enum damacy_numa_strategy)numa; + config.enable_gds = (enum damacy_gds_mode)gds; + struct damacy_reader* reader = component_value(reader_object, READER); + if (!reader) + return NULL; + struct damacy_executor* executor = NULL; + enum damacy_status status = + damacy_cuda_executor_create(reader, &config, &executor); + if (status != DAMACY_OK) + return api_raise_status(status, "create CUDA executor"); + return component_new(EXECUTOR, executor, reader_object); +} + +static PyObject* +compose_pipeline(PyObject* self, PyObject* args) +{ + (void)self; + PyObject *planner_object, *executor_object, *shape_object; + unsigned int dtype; + struct damacy_batch_spec output = { 0 }; + struct damacy_queue_limits queues; + if (!PyArg_ParseTuple(args, + "OOOIIII", + &planner_object, + &executor_object, + &shape_object, + &output.samples_per_batch, + &dtype, + &queues.lookahead_samples, + &queues.prepared_batches)) + return NULL; + output.dtype = (enum damacy_dtype)dtype; + struct damacy_planner* planner = component_value(planner_object, PLANNER); + if (!planner) + return NULL; + struct damacy_executor* executor = component_value(executor_object, EXECUTOR); + if (!executor) + return NULL; + PyObject* shape = PySequence_Fast(shape_object, "shape must be a sequence"); + if (!shape) + return NULL; + Py_ssize_t rank = PySequence_Fast_GET_SIZE(shape); + if (rank < 1 || rank > DAMACY_MAX_RANK) { + Py_DECREF(shape); + PyErr_SetString(PyExc_ValueError, "invalid output rank"); + return NULL; + } + output.sample_rank = (uint8_t)rank; + for (Py_ssize_t i = 0; i < rank; ++i) { + output.sample_shape[i] = + PyLong_AsLongLong(PySequence_Fast_GET_ITEM(shape, i)); + if (PyErr_Occurred()) { + Py_DECREF(shape); + return NULL; + } + } + Py_DECREF(shape); + PyObject* dependencies = PyTuple_Pack(2, planner_object, executor_object); + if (!dependencies) + return NULL; + PyObject* result = api_pipeline_from_components( + planner, executor, &output, &queues, dependencies); + Py_DECREF(dependencies); + return result; +} + +static PyMethodDef methods[] = { + { "create_reader", create_reader, METH_VARARGS, NULL }, + { "create_metadata_reader", create_metadata_reader, METH_VARARGS, NULL }, + { "create_metadata", create_metadata, METH_VARARGS, NULL }, + { "create_planner", create_planner, METH_VARARGS, NULL }, + { "create_cpu_executor", create_cpu_executor, METH_VARARGS, NULL }, + { "create_cuda_executor", create_cuda_executor, METH_VARARGS, NULL }, + { "compose_pipeline", compose_pipeline, METH_VARARGS, NULL }, + { NULL, NULL, 0, NULL } +}; + +int +components_register(PyObject* module) +{ + return PyModule_AddFunctions(module, methods); +} diff --git a/python/damacy/_native.c b/python/damacy/_native.c index 43140a50..7f62e524 100644 --- a/python/damacy/_native.c +++ b/python/damacy/_native.c @@ -13,7 +13,9 @@ #define PY_SSIZE_T_CLEAN #include +#ifdef DAMACY_HAS_CUDA #include +#endif #include "damacy.h" #include "damacy_limits.h" @@ -126,6 +128,7 @@ py_cuda_device_count(PyObject* self, PyObject* args) { (void)self; (void)args; +#ifdef DAMACY_HAS_CUDA CUresult r; int count = 0; if ((r = cuInit(0)) != CUDA_SUCCESS) @@ -133,6 +136,9 @@ py_cuda_device_count(PyObject* self, PyObject* args) if ((r = cuDeviceGetCount(&count)) != CUDA_SUCCESS) return PyLong_FromLong(0); return PyLong_FromLong((long)count); +#else + return PyLong_FromLong(0); +#endif } static PyObject* @@ -156,6 +162,7 @@ py_cuda_init_primary(PyObject* self, PyObject* args, PyObject* kw) if (!PyArg_ParseTupleAndKeywords(args, kw, "|i", kws, &device)) return NULL; +#ifdef DAMACY_HAS_CUDA CUresult r; if ((r = cuInit(0)) != CUDA_SUCCESS) { PyErr_Format(PyExc_RuntimeError, "cuInit failed (%d)", (int)r); @@ -178,6 +185,10 @@ py_cuda_init_primary(PyObject* self, PyObject* args, PyObject* kw) return NULL; } Py_RETURN_NONE; +#else + PyErr_SetString(PyExc_RuntimeError, "CUDA support is disabled in this build"); + return NULL; +#endif } static PyMethodDef methods[] = { @@ -247,8 +258,15 @@ module_exec(PyObject* m) PyErr_SetString(PyExc_RuntimeError, "failed to install damacy log sink"); return -1; } - if (api_register_types(m) != 0) + if (api_register_types(m) != 0 || components_register(m) != 0) return -1; +#ifdef DAMACY_HAS_CUDA + if (PyModule_AddIntConstant(m, "CUDA_ENABLED", 1) < 0) + return -1; +#else + if (PyModule_AddIntConstant(m, "CUDA_ENABLED", 0) < 0) + return -1; +#endif return 0; } diff --git a/python/damacy/_native.pyi b/python/damacy/_native.pyi index 83ae90a7..423487ef 100644 --- a/python/damacy/_native.pyi +++ b/python/damacy/_native.pyi @@ -6,6 +6,7 @@ from __future__ import annotations from typing import Any, Final __version__: Final[str] +CUDA_ENABLED: Final[int] # ---- log-level constants (mirror damacy_log.h) -------------------------- @@ -107,11 +108,11 @@ class Batch: - ``(1, 0)`` or higher → v1.0 ``"dltensor_versioned"`` capsule. CuPy and array-API-spec consumers ask for this. - ``dl_device`` is accepted for protocol compatibility but ignored; - the producer always emits on the assembling device.""" + ``dl_device`` must match the batch device. CPU tensors require + ``stream=None``.""" def __dlpack_device__(self) -> tuple[int, int]: - """Returns (kDLCUDA=2, ordinal).""" + """Returns (kDLCPU=1, 0) or (kDLCUDA=2, ordinal).""" class Pipeline: """Native streaming-pipeline handle. @@ -169,3 +170,74 @@ class Pipeline: def stats(self) -> dict[str, Any]: ... def stats_reset(self) -> None: ... + def shutdown(self) -> None: ... + +DEFAULT_CHUNK_UNCOMPRESSED_BYTES: Final[int] +DEFAULT_READ_OP_MAX_BYTES: Final[int] +DEFAULT_HOST_BUFFER_WAVES: Final[int] +DEFAULT_MAX_CHUNKS_PER_WAVE: Final[int] +DEFAULT_MAX_SUBSTREAMS_PER_CHUNK: Final[int] +DEFAULT_METADATA_IO_CONCURRENCY: Final[int] +DEFAULT_IO_THREADS: Final[int] +DEFAULT_ARRAY_META_CACHE: Final[int] +DEFAULT_SHARD_INDEX_CACHE: Final[int] +DEFAULT_CHUNK_LAYOUT_CACHE: Final[int] +DEFAULT_MAX_SHARDS_PER_SAMPLE: Final[int] +MAX_CHUNK_BYTES: Final[int] +MAX_READ_OP_BYTES: Final[int] +N_WAVES: Final[int] +MAX_HOST_BUFFER_WAVES: Final[int] +HARD_MAX_CHUNKS_PER_WAVE: Final[int] +HARD_MAX_SUBSTREAMS_PER_CHUNK: Final[int] +MAX_METADATA_IO_CONCURRENCY: Final[int] +MAX_IO_THREADS: Final[int] + +def create_reader(workers: int, max_inflight_reads: int, /) -> object: ... +def create_metadata_reader( + concurrency: int, + baseline_ns: int, + mu: float, + sigma: float, + cap_ns: int, + seed: int, + /, +) -> object: ... +def create_metadata( + reader: object, array_entries: int, shard_entries: int, / +) -> object: ... +def create_planner( + metadata: object, + max_chunks: int, + max_chunk_bytes: int, + max_shards: int, + max_plan_bytes: int, + /, +) -> object: ... +def create_cpu_executor( + reader: object, workers: int, max_encoded: int, max_decoded: int, max_memory: int, / +) -> object: ... +def create_cuda_executor( + reader: object, + device: int, + max_memory: int, + max_chunk: int, + max_read: int, + max_chunks: int, + max_substreams: int, + host_waves: int, + chunk_layout_entries: int, + numa_strategy: int, + numa_node: int, + gds: int, + /, +) -> object: ... +def compose_pipeline( + planner: object, + executor: object, + shape: tuple[int, ...], + samples: int, + dtype: int, + lookahead: int, + prepared_batches: int, + /, +) -> Pipeline: ... diff --git a/python/tests/test_components.py b/python/tests/test_components.py new file mode 100644 index 00000000..2877f028 --- /dev/null +++ b/python/tests/test_components.py @@ -0,0 +1,376 @@ +from __future__ import annotations + +import ctypes +import dataclasses +import gc +import json +import subprocess +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Barrier + +import damacy +import numpy as np +import pytest + + +def planner(**limits): + return damacy.ChunkPlanner( + metadata=damacy.ZarrMetadata( + reader=damacy.FileMetadataReader(concurrency=2), + cache=damacy.MetadataCache(array_entries=16, shard_index_entries=64), + ), + limits=damacy.PlanLimits( + **( + {"max_chunks": 128, "max_chunk_bytes": 1024, "max_shards_per_sample": 4} + | limits + ) + ), + ) + + +def executor(**limits): + return damacy.CpuExecutor( + reader=damacy.FileReader(workers=2, max_inflight_reads=1), + limits=damacy.CpuLimits( + **( + { + "decode_workers": 2, + "max_encoded_chunk_bytes": 1024, + "max_decoded_chunk_bytes": 1024, + "max_memory_bytes": 8 << 20, + } + | limits + ) + ), + ) + + +def pipeline(*, shape=(8, 16), samples=1, **overrides): + return damacy.Pipeline( + **( + { + "planner": planner(), + "executor": executor(), + "output": damacy.BatchSpec(samples=samples, shape=shape), + "queues": damacy.QueueLimits(lookahead_samples=4), + "pop_timeout_s": 5.0, + } + | overrides + ) + ) + + +def sample(uri, y=0, x=0, shape=(8, 16)): + return damacy.Sample(uri=uri, aabb=[(y, y + shape[0]), (x, x + shape[1])]) + + +def raw_array(path: Path, values, *, codec=None, payload=None, fill=0): + values = np.asarray(values) + codecs = [{"name": "bytes", "configuration": {"endian": "little"}}] + if codec is not None: + codecs.append(codec) + path.mkdir() + (path / "zarr.json").write_text( + json.dumps( + { + "zarr_format": 3, + "node_type": "array", + "shape": list(values.shape), + "data_type": values.dtype.name, + "fill_value": fill, + "chunk_grid": { + "name": "regular", + "configuration": {"chunk_shape": list(values.shape)}, + }, + "chunk_key_encoding": { + "name": "default", + "configuration": {"separator": "/"}, + }, + "codecs": codecs, + } + ) + ) + chunk = path / "c" / "0" / "0" + chunk.parent.mkdir(parents=True) + chunk.write_bytes(values.tobytes() if payload is None else payload) + return str(path) + + +def test_cpu_numpy_and_reader_backpressure(tiny_zarr): + expected = np.arange(128, dtype=np.float32).reshape(8, 16)[1:7, 3:12] + with pipeline(shape=(6, 9), samples=2) as p: + p.push(sample(tiny_zarr, 1, 3, (6, 9)) for _ in range(12)) + for i in range(6): + with p.pop() as batch: + assert batch.info.batch_id == i + assert batch.info.device_type == damacy.DeviceType.CPU + assert batch.info.device_id == 0 + assert batch.info.ready_stream == 0 + assert batch.__dlpack_device__() == (1, 0) + view = np.from_dlpack(batch) + assert view.ctypes.data == batch.info.data + assert view.dtype == np.float32 + np.testing.assert_array_equal(view, np.stack([expected, expected])) + del view + stats = p.stats() + assert stats.gpu_bytes_committed == 0 + assert 0 < stats.host_bytes_committed <= 8 << 20 + assert stats.chunks_dispatched < stats.chunks_planned + assert p.device == -1 + + +def test_retained_numpy_views_and_shutdown(tiny_zarr): + p = pipeline(shape=(2, 3), pop_timeout_s=0.05) + p.push( + [ + sample(tiny_zarr, 0, 0, (2, 3)), + sample(tiny_zarr, 2, 4, (2, 3)), + sample(tiny_zarr, 6, 10, (2, 3)), + ] + ) + first, second = p.pop(), p.pop() + first_view = np.from_dlpack(first) + expected = first_view.copy() + first.release() + with pytest.raises(damacy.PoolStarved): + p.pop() + np.testing.assert_array_equal(first_view, expected) + second.release() + third = p.pop() + third_view = np.from_dlpack(third) + third_expected = np.arange(128, dtype=np.float32).reshape(8, 16)[6:8, 10:13][None] + third.release() + p.close() + del first, second, third, p + gc.collect() + np.testing.assert_array_equal(first_view, expected) + np.testing.assert_array_equal(third_view, third_expected) + + +def test_close_wakes_timed_out_pop(tiny_zarr): + p = pipeline(samples=2, pop_timeout_s=0.02) + p.push([sample(tiny_zarr)]) + with pytest.raises(damacy.PoolStarved): + p.pop() + worker = p._pop_thread + start = time.monotonic() + p.close() + assert time.monotonic() - start < 2 + assert worker is not None and not worker.is_alive() + p.close() + with pytest.raises(damacy.ShutdownError): + p.pop() + + +def test_components_exclusive_and_reusable(tiny_zarr): + components = {"planner": planner(), "executor": executor()} + retained = [] + for _ in range(2): + with pipeline(**components) as p: + with pytest.raises(damacy.InvalidArgument): + pipeline(**components) + p.push([sample(tiny_zarr)]) + with p.pop() as batch: + retained.append(np.from_dlpack(batch)) + for view in retained: + np.testing.assert_array_equal( + view, np.arange(128, dtype=np.float32).reshape(1, 8, 16) + ) + + +def test_mixed_source_types(tmp_path): + a = np.arange(16, dtype=np.uint16).reshape(4, 4) + b = (np.arange(16, dtype=np.float32) / 8 - 2).reshape(4, 4) + uris = [raw_array(tmp_path / "a", a), raw_array(tmp_path / "b", b)] + with pipeline(shape=(4, 4), samples=2) as p: + p.push(sample(uri, shape=(4, 4)) for uri in uris) + with p.pop() as batch: + np.testing.assert_array_equal(np.from_dlpack(batch), np.stack([a, b])) + + +@pytest.mark.parametrize( + "codec,payload", + [ + (None, b"short"), + ( + {"name": "zstd", "configuration": {"level": 3, "checksum": False}}, + b"not a zstd frame", + ), + ( + { + "name": "blosc", + "configuration": { + "cname": "zstd", + "clevel": 3, + "shuffle": "shuffle", + "typesize": 4, + "blocksize": 0, + }, + }, + bytes(16), + ), + ], +) +def test_corrupt_chunks_fail_terminally(tmp_path, codec, payload): + uri = raw_array( + tmp_path / "bad", + np.arange(16, dtype=np.float32).reshape(4, 4), + codec=codec, + payload=payload, + ) + with pipeline(shape=(4, 4)) as p: + p.push([sample(uri, shape=(4, 4))]) + for _ in range(2): + with pytest.raises(damacy.DecodeError): + p.pop() + + +def test_memory_and_plan_limits(tiny_zarr): + with pytest.raises(damacy.BudgetExceeded): + pipeline(executor=executor(max_memory_bytes=4096)) + with pipeline(planner=planner(max_plan_bytes=1)) as p: + p.push([sample(tiny_zarr)]) + with pytest.raises(damacy.BudgetExceeded): + p.pop() + with pipeline(executor=executor(max_decoded_chunk_bytes=8)) as p: + p.push([sample(tiny_zarr)]) + with pytest.raises(damacy.BudgetExceeded): + p.pop() + + +def test_cpu_dlpack_parameters(tiny_zarr): + with pipeline() as p: + p.push([sample(tiny_zarr)]) + with p.pop() as batch: + with pytest.raises(ValueError, match="stream=None"): + batch.__dlpack__(stream=1) + with pytest.raises(BufferError, match="device"): + batch.__dlpack__(dl_device=(2, 0)) + with pytest.raises(BufferError, match="copy=True"): + batch.__dlpack__(copy=True) + cap = batch.__dlpack__(dl_device=(1, 0), max_version=(1, 0)) + del cap + + +def test_invalid_limits(): + for make in [ + lambda: damacy.BatchSpec(0, (4,)), + lambda: damacy.BatchSpec(1, (0,)), + lambda: damacy.QueueLimits(0), + lambda: damacy.QueueLimits(4, 0), + lambda: damacy.CpuLimits(0), + lambda: damacy.PlanLimits(max_chunks=0), + lambda: damacy.CpuLimits(1 << 20, decode_workers=-1), + ]: + with pytest.raises(ValueError): + make() + with pytest.raises(TypeError): + damacy.Pipeline(planner=planner()) + assert dataclasses.replace(damacy.BatchSpec(1, (4,)), samples=2).samples == 2 + + +@pytest.mark.parametrize("dtype", ["f32", "bf16"]) +def test_cpu_torch_consumer_retains_tensor(tiny_zarr, dtype): + torch = pytest.importorskip("torch") + with pipeline(output=damacy.BatchSpec(1, (8, 16), dtype)) as p: + p.push([sample(tiny_zarr)]) + with p.pop() as batch: + tensor = torch.from_dlpack(batch) + assert tensor.device.type == "cpu" + assert tensor.dtype == (torch.float32 if dtype == "f32" else torch.bfloat16) + del batch + expected = torch.arange(128, dtype=torch.float32).reshape(1, 8, 16) + assert torch.equal(tensor.float(), expected) + + +@pytest.mark.usefixtures("cuda_ctx") +@pytest.mark.parametrize("read_capacity", [1, 16]) +def test_cuda_composition_matches_cpu(tiny_zarr, read_capacity): + cuda = damacy.CudaExecutor( + reader=damacy.FileReader(workers=2, max_inflight_reads=read_capacity), + device=0, + limits=damacy.CudaLimits( + max_gpu_memory_bytes=1 << 30, + max_chunk_bytes=1024, + max_chunks_per_wave=16, + max_substreams_per_chunk=16, + ), + ) + driver = ctypes.CDLL("libcuda.so.1") + copy = driver.cuMemcpyDtoH_v2 + copy.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_size_t] + copy.restype = ctypes.c_int + with pipeline(shape=(6, 9), executor=cuda) as p: + p.push([sample(tiny_zarr, 1, 3, (6, 9))]) + batch = p.pop() + assert batch.__dlpack_device__() == (2, 0) + result = np.empty(batch.info.shape, dtype=np.float32) + assert copy(result.ctypes.data, batch.info.data, result.nbytes) == 0 + assert copy(result.ctypes.data, batch.info.data, result.nbytes) == 0 + capsule = batch.__dlpack__(stream=None) + batch.release() + del capsule + expected = np.arange(128, dtype=np.float32).reshape(8, 16)[1:7, 3:12] + np.testing.assert_array_equal(result, expected[None]) + with pipeline(shape=(6, 9)) as p: + p.push([sample(tiny_zarr, 1, 3, (6, 9))]) + with p.pop() as batch: + np.testing.assert_array_equal(result, np.from_dlpack(batch)) + + +@pytest.mark.parametrize("shuffle", ["noshuffle", "shuffle", "bitshuffle"]) +def test_blosc_filters(tmp_path, write_zarr_script, shuffle): + uri = tmp_path / "array" + subprocess.run( + [ + "uv", + "run", + "--script", + str(write_zarr_script), + "--out", + str(uri), + "--shape", + "8,16", + "--inner", + "4,8", + "--shard", + "8,16", + "--dtype", + "float32", + "--codec", + "blosc-zstd", + "--shuffle", + shuffle, + ], + check=True, + capture_output=True, + ) + with pipeline(shape=(6, 9)) as p: + p.push([sample(str(uri), 1, 3, (6, 9))]) + with p.pop() as batch: + expected = np.arange(128, dtype=np.float32).reshape(8, 16)[1:7, 3:12] + np.testing.assert_array_equal(np.from_dlpack(batch), expected[None]) + + +def test_parallel_construction_claims_components_once(): + components = {"planner": planner(), "executor": executor()} + barrier = Barrier(2) + + def create(): + barrier.wait() + try: + return pipeline(**components) + except damacy.InvalidArgument: + return None + + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(create) for _ in range(2)] + results = [future.result(timeout=5) for future in futures] + active = [result for result in results if result is not None] + try: + assert len(active) == 1 + finally: + for result in active: + result.close() diff --git a/python/tests/test_damacy.py b/python/tests/test_damacy.py index 1b17258e..b75ae596 100644 --- a/python/tests/test_damacy.py +++ b/python/tests/test_damacy.py @@ -13,7 +13,6 @@ import dataclasses import shutil import subprocess -import sys import warnings import damacy @@ -256,12 +255,14 @@ def test_dtype_string_form_accepted(tiny_zarr, dtype): _ = tiny_zarr with Pipeline(_base_config(dtype=dtype)) as d: assert isinstance(d, Pipeline) + assert d.config is not None assert d.config.dtype is damacy.Dtype.coerce(dtype) def test_dtype_int_form_accepted(tiny_zarr): _ = tiny_zarr with Pipeline(_base_config(dtype=damacy.Dtype.BF16)) as d: + assert d.config is not None assert d.config.dtype is damacy.Dtype.BF16 @@ -614,9 +615,7 @@ def test_config_validates_eagerly(): def test_config_tuning_defaults_are_explicit(): cfg = Config(samples_per_batch=1, sample_shape=(8, 16), max_gpu_memory_bytes=1) - assert ( - cfg.max_chunk_uncompressed_bytes == _native.DEFAULT_CHUNK_UNCOMPRESSED_BYTES - ) + assert cfg.max_chunk_uncompressed_bytes == _native.DEFAULT_CHUNK_UNCOMPRESSED_BYTES assert cfg.max_read_op_bytes == _native.DEFAULT_READ_OP_MAX_BYTES assert cfg.host_buffer_waves == _native.DEFAULT_HOST_BUFFER_WAVES assert cfg.max_chunks_per_wave == _native.DEFAULT_MAX_CHUNKS_PER_WAVE @@ -1052,26 +1051,36 @@ def test_batch_dlpack_after_release_raises(tiny_zarr): batch.__dlpack_device__() -def test_batch_dlpack_capsule_holds_batch_alive(tiny_zarr): - """The capsule's Py_INCREF lands on the C-side BatchObj - (``batch._native``), not on the Python-level wrapper — the wrapper - just delegates. Dropping the capsule must run the deleter and - decrement that native reference back to its baseline.""" - uri = tiny_zarr - with Pipeline(_base_config()) as d: - d.push([Sample(uri=uri, aabb=[(0, 8), (0, 16)])]) +@pytest.mark.parametrize("version", [(0, 8), (1, 0)]) +def test_batch_dlpack_capsule_retains_data_after_shutdown(tiny_zarr, version): + driver = ctypes.CDLL("libcuda.so.1") + copy = driver.cuMemcpyDtoH_v2 + copy.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_size_t] + copy.restype = ctypes.c_int + config = dataclasses.replace(_base_config(), sample_shape=(4, 8)) + with Pipeline(config) as d: + d.push([Sample(uri=tiny_zarr, aabb=[(0, 4), (0, 8)])]) batch = d.pop() - try: - native = batch._native - rc_before = sys.getrefcount(native) - cap = batch.__dlpack__(stream=None) - rc_held = sys.getrefcount(native) - assert rc_held > rc_before, "capsule should incref the native batch" - del cap - rc_after = sys.getrefcount(native) - assert rc_after == rc_before, "capsule deleter should decref it" - finally: - batch.release() + cap = batch.__dlpack__(stream=None, max_version=version) + name, layout = ( + (b"dltensor", _DLManagedTensor) + if version[0] == 0 + else (b"dltensor_versioned", _DLManagedTensorVersioned) + ) + tensor = _capsule_as(cap, name, layout).dl_tensor + before = (ctypes.c_float * 32)() + assert copy(before, tensor.data, ctypes.sizeof(before)) == 0 + batch.release() + for _ in range(3): + d.push([Sample(uri=tiny_zarr, aabb=[(4, 8), (8, 16)])]) + with d.pop(): + pass + del d, batch + after = (ctypes.c_float * 32)() + assert copy(after, tensor.data, ctypes.sizeof(after)) == 0 + assert list(after) == list(before) + assert [tensor.shape[i] for i in range(tensor.ndim)] == [1, 4, 8] + del cap def test_batch_dlpack_stream_kwargs_accepted(tiny_zarr): diff --git a/python/tests/test_deferred_release.py b/python/tests/test_deferred_release.py index 527282b9..9bfc3dc2 100644 --- a/python/tests/test_deferred_release.py +++ b/python/tests/test_deferred_release.py @@ -58,6 +58,7 @@ def _mk_cfg() -> Config: ) +@pytest.mark.usefixtures("cuda_ctx") def test_release_event_none_falls_back_to_immediate(one_zarr): """``release(event=None)`` is the same as ``release()``.""" with Pipeline(_mk_cfg()) as d: @@ -67,6 +68,7 @@ def test_release_event_none_falls_back_to_immediate(one_zarr): b.release() # idempotent +@pytest.mark.usefixtures("cuda_ctx") def test_release_event_rejects_bad_type(one_zarr): with Pipeline(_mk_cfg()) as d: d.push([Sample(uri=one_zarr, aabb=[(0, 8), (0, 16)])]) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 49ab2d1e..31af977c 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -13,7 +13,7 @@ add_src_lib(log SOURCES log/log.h log/log.c ${PROJECT_SOURCE_DIR}/src/damacy_log # in libFuzzer). The .c is always in the source list (gives the static # lib a linker language); its body is #if-gated on DAMACY_NVTX_ENABLED. add_cuda_lib(damacy_nvtx SOURCES nvtx/nvtx.h nvtx/nvtx.c LINKS CUDA::cuda_driver platform) -if(NOT DAMACY_FUZZ AND DAMACY_NVTX_ENABLED) +if(DAMACY_CUDA AND DAMACY_NVTX_ENABLED) target_link_libraries(damacy_nvtx PUBLIC CUDA::nvtx3) target_compile_definitions(damacy_nvtx PUBLIC DAMACY_NVTX_ENABLED=1) endif() @@ -45,7 +45,10 @@ add_platform_sources(platform platform/numa) # CUDA-side GPU→host-NUMA-node resolution + scope/affinity wrappers over # the platform NUMA primitives. CUDA::cuda_driver for cuDeviceGetAttribute # (HOST_NUMA_ID) and the BDF sysfs fallback. -add_cuda_lib(numa SOURCES numa/numa.h numa/numa.c LINKS log platform CUDA::cuda_driver) +add_src_lib(numa SOURCES numa/numa.h numa/affinity.c LINKS log platform) +add_cuda_lib(numa_cuda SOURCES numa/numa_cuda.h numa/numa.c + LINKS numa log platform CUDA::cuda_driver +) add_src_lib(platform_io SOURCES platform/platform_io.h LINKS platform) add_platform_sources(platform_io platform/platform_io) @@ -112,13 +115,13 @@ add_cuda_lib(gpu_budget SOURCES gpu_budget/gpu_budget.h gpu_budget/gpu_budget.c # Two-slot batch state machine + lazy device-tensor pool. Owns the # caller-visible minibatch lifetime; render_job owns planner scratch. add_cuda_lib(batch_pool SOURCES batch_pool/batch_pool.h batch_pool/batch_pool.c - LINKS planner path_intern log CUDA::cuda_driver + LINKS dispatch_utils path_intern log CUDA::cuda_driver ) # Planner output plus dispatch cursor for one sealed batch. V1 keeps # one render job per batch slot. add_cuda_lib(render_job SOURCES render_job/render_job.h render_job/render_job.c - LINKS planner path_intern store log CUDA::cuda_driver + LINKS dispatch_utils path_intern store log CUDA::cuda_driver ) # Input staging slot pool. @@ -143,7 +146,7 @@ add_cuda_lib(fanout SOURCES wave/fanout.h wave/fanout.c # cuStreamSynchronize. add_cuda_lib(wave_budget SOURCES wave/input_transfer.h wave/wave_budget.h wave/wave_budget.c - LINKS input_transfer decoder fanout planner assemble gpu_budget + LINKS input_transfer decoder fanout dispatch_utils assemble gpu_budget damacy_config log CUDA::cuda_driver ) @@ -164,7 +167,7 @@ add_cuda_lib(wave SOURCES wave/wave.h wave/wave.c add_cuda_lib(wave_pool SOURCES wave/input_transfer.h wave/wave_input.h wave/wave_input.c wave/wave_pool.h wave/wave_pool.c LINKS wave input_slot input_transfer fanout wave_budget gpu_budget - batch_pool decoder assemble planner render_job store damacy_stats + batch_pool decoder assemble dispatch_utils render_job store damacy_stats damacy_config damacy_nvtx log CUDA::cuda_driver ) @@ -177,7 +180,7 @@ add_src_lib(store store/store_latency.h store/store_latency.c LINKS io_queue platform_io strbuf log lru pool hash m ) -if(NOT DAMACY_FUZZ AND DAMACY_ENABLE_GDS) +if(DAMACY_CUDA AND DAMACY_ENABLE_GDS) target_sources(store PRIVATE store/store_fs_gds.h store/store_fs_gds.c) target_include_directories(store PRIVATE ${CuFile_INCLUDE_DIR}) target_link_libraries(store PRIVATE CUDA::cuda_driver ${CMAKE_DL_LIBS}) @@ -214,16 +217,15 @@ add_src_lib(zarr_chunk_layout SOURCES zarr/zarr_chunk_layout.h zarr/zarr_chunk_l LINKS zarr store log ) -# Planner: samples + cached metadata → page-aligned read_ops + chunk_plans. -# coalesce + group_chunks are the post-emit steps of the IO planning -# pipeline (filter → sort → fuse-with-cap → group-by-read). -add_src_lib(planner - SOURCES planner/planner.h planner/planner.c - planner/coalesce.h planner/coalesce.c - planner/group_chunks.h planner/group_chunks.c - planner/read_op_sort.h planner/read_op_sort.c - LINKS prefetch_cache shard_index zarr dtype platform strbuf - path_intern log damacy_config +add_src_lib(dispatch_utils + SOURCES executor/coalesce.h executor/coalesce.c + executor/group_chunks.h executor/group_chunks.c + executor/read_op_sort.h executor/read_op_sort.c + LINKS path_intern log damacy_config +) + +add_src_lib(planner SOURCES planner/planner.h planner/planner.c + LINKS plan_builder execution_plan prefetch_cache log ) # GPU pieces — add_cuda_lib is a no-op under DAMACY_FUZZ. @@ -237,7 +239,7 @@ add_cuda_lib(decoder decoder/launch_check.h decoder/launch_check.cu decoder/status_reduce.h decoder/status_reduce.cu LINKS CUDA::cudart_static CUDA::cuda_driver nvcomp::nvcomp log zarr - planner assemble zarr_chunk_layout + dispatch_utils assemble zarr_chunk_layout ) # Bytewise gather kernel: chunk windows → output batch tensor. @@ -248,19 +250,44 @@ add_cuda_lib(assemble SOURCES assemble/assemble.h assemble/assemble.cu LINKS CUDA::cudart_static CUDA::cuda_driver log ) -# Public damacy surface. The API implementation uses the CUDA driver API -# (cuMem*, cuStream*, cuEvent*); the runtime API is only pulled in -# transitively via decoder/assemble (nvcomp's headers + the kernel -# launch in assemble.cu). -add_cuda_lib(damacy - SOURCES damacy.h damacy_internal.h - damacy_lifecycle.c damacy_push.c damacy_plan.c damacy_pop.c - damacy_scheduler.c damacy_status.c - LINKS platform planner zarr store - decoder assemble dtype log lookahead gpu_budget - batch_pool render_job wave_pool wave_budget damacy_stats damacy_config - damacy_nvtx strbuf scheduler io_queue metadata_store_async - prefetch_cache prefetcher - array_meta shard_index chunk_layout - CUDA::cuda_driver +add_src_lib(prepared_plan SOURCES planner/plan.h planner/plan.c) + +add_src_lib(plan_builder SOURCES planner/plan_builder.h planner/plan_builder.c + LINKS prepared_plan prefetch_cache shard_index zarr damacy_config strbuf hash +) + +add_src_lib(execution_plan SOURCES executor/dispatch.h executor/dispatch.c + LINKS dispatch_utils prepared_plan path_intern ) + +if(NOT DAMACY_FUZZ) + add_src_lib(pipeline_components SOURCES pipeline/components.h pipeline/components.c damacy_status.c + LINKS damacy_config store platform log + ) + add_src_lib(metadata_planner SOURCES pipeline/zarr_planner.c + LINKS pipeline_components plan_builder prefetcher array_meta shard_index + metadata_store_async lookahead + ) + add_src_lib(cpu_executor SOURCES executor/cpu_executor.c + LINKS pipeline_components prepared_plan threadpool damacy_stats + PkgConfig::ZSTD PkgConfig::BLOSC + ) + add_src_lib(damacy + SOURCES damacy.h damacy_pipeline.h damacy_internal.h + damacy_lifecycle.c damacy_push.c damacy_plan.c damacy_pop.c + damacy_scheduler.c + LINKS pipeline_components metadata_planner cpu_executor scheduler damacy_stats + ) + if(DAMACY_CUDA) + add_cuda_lib(cuda_executor + SOURCES executor/cuda_executor.h executor/cuda_executor.c executor/cuda_geometry.inc + LINKS pipeline_components execution_plan numa_cuda gpu_budget batch_pool + render_job wave_pool wave_budget damacy_config zarr_chunk_layout + CUDA::cuda_driver + ) + target_link_libraries(damacy PUBLIC cuda_executor) + target_compile_definitions(damacy PUBLIC DAMACY_HAS_CUDA=1) + else() + target_sources(damacy PRIVATE executor/cuda_stub.c) + endif() +endif() diff --git a/src/assemble/assemble.h b/src/assemble/assemble.h index 153c64a4..880320c8 100644 --- a/src/assemble/assemble.h +++ b/src/assemble/assemble.h @@ -12,7 +12,7 @@ #include "damacy.h" // enum damacy_dtype (destination) #include "damacy_limits.h" -#include "planner/planner.h" +#include "executor/dispatch.h" #include #include diff --git a/src/batch_pool/batch_pool.c b/src/batch_pool/batch_pool.c index 462d9e6d..cd8d5007 100644 --- a/src/batch_pool/batch_pool.c +++ b/src/batch_pool/batch_pool.c @@ -1,7 +1,8 @@ #include "batch_pool.h" +#include "executor/dispatch.h" #include "log/log.h" -#include "planner/planner.h" +#include "planner/plan_builder.h" #include "util/cuda_check.h" // CU + CUDPTR #include "util/prelude.h" diff --git a/src/damacy.h b/src/damacy.h index 4ca3267c..077c1882 100644 --- a/src/damacy.h +++ b/src/damacy.h @@ -188,6 +188,16 @@ extern "C" struct damacy; struct damacy_batch; + enum damacy_device_type + { + DAMACY_DEVICE_CPU = 1, + DAMACY_DEVICE_CUDA = 2, + }; + + void damacy_batch_retain(struct damacy_batch* batch); + void damacy_batch_release(struct damacy_batch* batch); + void damacy_shutdown(struct damacy* d); + // Fill performance/resource knobs with explicit library defaults. Callers // still own required geometry fields such as sample_shape, samples_per_batch, // lookahead_samples, dtype, and max_gpu_memory_bytes. @@ -204,7 +214,7 @@ extern "C" // Requires a live CUDA context on the calling thread. void damacy_config_describe(const struct damacy_config* cfg); - // The CUDA device index this instance is bound to. + // CUDA device index, or -1 for CPU execution. int damacy_get_device(const struct damacy* d); // Tear down. Does NOT flush in-flight work; the io_queue is asked to @@ -234,50 +244,29 @@ extern "C" struct damacy_push_result damacy_push(struct damacy* d, struct damacy_sample_slice samples); - // Block until the next batch is on-device-ready, in push-FIFO order. - // *out is owned by damacy until damacy_release. + // Return the next ready batch in push order. The caller owns one reference; + // release it with damacy_release or damacy_batch_release. enum damacy_status damacy_pop(struct damacy* d, struct damacy_batch** out); - // Return the batch's slot to the pool. Thread-safe; may be called from + // Release one batch reference. The buffer is reusable after the last + // consumer releases it. Thread-safe; may be called from // a thread other than the one that called damacy_pop. void damacy_release(struct damacy* d, struct damacy_batch* b); - // Deferred release: tell damacy not to reuse the batch's buffer until - // `event` (a CUevent) has fired. Useful when the consumer kicked off - // an async D2D copy on a side stream and wants the host to return - // immediately instead of blocking on cuEventSynchronize before exiting - // a `with` block. - // - // Damacy records the wait on its internal stream_post (which is where - // assemble writes the slot's output buffer), so the next batch's - // assemble kernel — which writes to the same buffer — will wait on - // `event` before launching. No host synchronization is performed; this - // call returns as soon as the wait is queued and the slot state - // transitions to FREE. - // - // Latency note: stream_post is shared across both batch slots, so the - // wait gates EVERY subsequent assemble until `event` fires — not just - // the released slot's next reuse. A long-held consumer event therefore - // stalls the second slot's assemble too. For maximum overlap, release - // with an event that completes quickly relative to the consumer's - // step time. - // - // The event must remain valid for the duration of this call; damacy - // captures it into stream_post's command queue via cuStreamWaitEvent - // and is done with the handle by return. Passing a NULL event is - // equivalent to damacy_release. - // - // Returns DAMACY_OK on success or DAMACY_CUDA if the driver call fails. - // In either case the slot is released back to the pool — on the - // DAMACY_CUDA path the deferred wait was not installed and the slot - // falls back to immediate release, so the caller knows reuse is not - // gated on `event` but won't block on a future pop. + // Release one reference after ordering CUDA output writes behind event. + // While active, the wait is queued on the shared output stream and gates + // both output slots. After shutdown this synchronizes the event on the host. + // CPU execution rejects non-NULL events. The reference is released even if + // the wait fails; callers must handle that error before allowing reuse. enum damacy_status damacy_release_event(struct damacy* d, struct damacy_batch* b, void* event); struct damacy_batch_info { + void* data; + enum damacy_device_type device_type; + int device_id; void* device_ptr; // dtype-typed, contiguous int64_t shape[DAMACY_MAX_RANK + 1]; // [N, ...zarr axes] uint8_t rank; // includes leading N axis @@ -308,7 +297,8 @@ extern "C" struct damacy_metric plan; struct damacy_metric io; struct damacy_metric input_transfer; - // decode: stream_decode work only (nvcomp + status_reduce). + // CPU decode/assemble times sum worker observations. CUDA decode measures + // stream_decode work (nvcomp + status_reduce). // post_decode: stream_post work — post-decode kernels + 4B D2H + // cross-stream wait on decode_done. A large post_decode avg means // stream_post is bottlenecking; a small avg means it overlaps the @@ -382,6 +372,7 @@ extern "C" // wave-init to the first damacy_pop (lazy batch pool sizing) and stays // flat after that. Useful for surfacing the runtime budget to callers. uint64_t gpu_bytes_committed; + uint64_t host_bytes_committed; }; void damacy_stats_get(const struct damacy* d, struct damacy_stats* out); diff --git a/src/damacy_internal.h b/src/damacy_internal.h index 9b154f2a..09fb80bf 100644 --- a/src/damacy_internal.h +++ b/src/damacy_internal.h @@ -1,130 +1,36 @@ #pragma once -#include "batch_pool/batch_pool.h" -#include "damacy.h" -#include "gpu_budget/gpu_budget.h" -#include "lookahead/lookahead.h" -#include "numa/numa.h" -#include "planner/planner.h" -#include "prefetch/array_meta.h" -#include "prefetch/chunk_layout.h" -#include "prefetch/prefetch_cache.h" -#include "prefetch/prefetcher.h" -#include "prefetch/shard_index.h" -#include "render_job/render_job.h" +#include "pipeline/components.h" #include "scheduler/scheduler.h" -#include "store/metadata_store_async.h" -#include "store/store.h" -#include "wave/wave_pool.h" -#include -#include - -// Only one handle is live at a time (the orchestrator's `handle` field). -struct damacy_batch -{ - struct damacy* d; - uint16_t slot_idx; - uint64_t batch_id; -}; - -// Lock order (outer → inner): scheduler_lock → prefetcher->lock → cache->lock -// → lookahead->lock. Code holding scheduler_lock may acquire any of the -// inner locks; the reverse causes deadlock. The prefetcher worker never -// takes scheduler_lock. struct damacy { - struct damacy_config cfg; - enum damacy_status failed_status; + struct damacy_batch_spec output; + struct damacy_queue_limits queues; + struct damacy_planner* planner; + struct damacy_executor* executor; + struct prepared_plan** plans; + uint32_t plan_head; + uint32_t plan_count; uint64_t next_batch_id; - // Push-side cursor: indexes into the stream of samples ever pushed. The - // prefetcher consumes the same global sample_seq order independently of - // output batch placement. - uint64_t pushed_samples; - uint64_t page_alignment; - int cuda_device; - int retained_primary_device; // -1 = caller's ctx; else release at destroy - CUcontext retained_primary; // pushed per-call by ctx_guard when retained - CUcontext worker_ctx; // pushed by the worker on its first tick - - // Resolved NUMA placement plan; node<0 means "no pinning". Filled by - // numa_init at create-time and shared with the store's io_queue and - // the scheduler so worker threads can pin themselves on entry. - struct numa_resolved numa; - - // GPU memory budgeting. Single source of truth for committed/max - // across wave_pool, batch_pool, and the observe-and-grow paths. - struct gpu_budget* budget; - - struct store* store_host; - struct metadata_store_async* store_meta_async; - struct store* store_gds; - struct planner* planner; - - // Async metadata cache fetchers use a separate store boundary so cache - // misses do not occupy the prefetcher worker while waiting on small I/O. - struct array_meta_async_fetcher array_meta_async_fetcher; - struct shard_index_async_fetcher shard_index_async_fetcher; - struct chunk_layout_async_fetcher chunk_layout_async_fetcher; - struct prefetch_cache* array_meta_cache; - struct prefetch_cache* shard_index_cache; - struct prefetch_cache* chunk_layout_cache; - struct prefetcher* prefetcher; - // plan_ready_prefetch pops ready sample_seq intervals here, then - // plan_reserve copies each segment into the target batch slot. - struct prefetcher_ready* staging; - - struct damacy_lookahead lookahead; - struct damacy_batch_pool batch_pool; - struct render_job_pool render_jobs; - // Owns the 4 streams + both waves; built once in damacy_create and - // driven directly by the orchestrator (no per-call ctx building). - struct wave_pool wave_pool; - - struct damacy_batch handle; + enum damacy_status failed_status; struct damacy_stats stats; - - // Worker drives the pipeline; user-thread API coordinates via scheduler_lock. - // Created last in damacy_create, torn down first in destroy_inner. struct scheduler* sched; - int worker_ctx_pushed; // worker pushes worker_ctx on first tick; no matching - // pop. -}; - -// Tight enough to react to GPU event completion before wave-boundary -// gaps dominate; cuEventQuery is cheap so the worker isn't burning a core. -#define DAMACY_POP_POLL_NS 10000 - -struct ctx_guard -{ - int active; + uint32_t pop_calls; + _Atomic int stopping; + _Atomic int stopped; + int device; + int owns_components; + struct damacy_reader* owned_reader; + struct damacy_metadata_reader* owned_metadata_reader; + struct damacy_metadata* owned_metadata; }; -enum damacy_status -ctx_guard_enter(struct damacy* d, struct ctx_guard* g); - -void -ctx_guard_exit(struct ctx_guard* g); - int damacy_scheduler_step(void* arg); - -enum damacy_status -plan_reserve(struct damacy* self, - uint16_t slot_idx, - struct prefetcher_ready* ready, - uint32_t n_ready, - int close_batch); - -enum damacy_status -plan_run(struct damacy* self, uint16_t slot_idx, float* out_elapsed_ms); - -enum damacy_status -plan_commit(struct damacy* self, - uint16_t slot_idx, - enum damacy_status run_status, - float elapsed_ms, - int* changed); - +void +damacy_scheduler_enter(void* arg); +void +damacy_scheduler_leave(void* arg); enum damacy_status -plan_ready_prefetch(struct damacy* self, int* changed); +pipeline_prepare(struct damacy* self, int* changed); diff --git a/src/damacy_lifecycle.c b/src/damacy_lifecycle.c index b6ffc04c..2fd8e5a4 100644 --- a/src/damacy_lifecycle.c +++ b/src/damacy_lifecycle.c @@ -1,483 +1,132 @@ -#include "damacy.h" +#include "damacy_internal.h" #include "damacy_config.h" -#include "damacy_internal.h" #include "damacy_stats.h" -#include "log/log.h" -#include "platform/platform.h" -#include "store/store_fs_gds.h" -#include "util/cuda_check.h" -#include "util/prelude.h" -#include "wave/wave_budget.h" - -#include - -// --- ctx guard ------------------------------------------------------------ - -enum damacy_status -ctx_guard_enter(struct damacy* d, struct ctx_guard* g) -{ - g->active = 0; - if (!d || d->retained_primary_device < 0) - return DAMACY_OK; - enum damacy_status s = DAMACY_CUDA; - CU(Fail, cuCtxPushCurrent(d->retained_primary)); - g->active = 1; - return DAMACY_OK; -Fail: - return s; -} - -void -ctx_guard_exit(struct ctx_guard* g) -{ - if (g && g->active) { - cuCtxPopCurrent(NULL); - g->active = 0; - } -} - -// --- public API: create / destroy ---------------------------------------- - -// Single teardown list shared by every destroy path. cuda_skip=1 leaks -// CUDA-owned state and skips driver calls; CPU heap is always released. -// wave_pool first: its destroy syncs streams before the downstream -// batch_pool / planner free what those streams touched. -static void -destroy_inner(struct damacy* self, int cuda_skip) -{ - if (!self) - return; - - // Stop the worker first so its accesses retire before we free. - scheduler_destroy(self->sched); - self->sched = NULL; - - // Stop the prefetcher thread first so no new metadata requests are admitted. - // Its gates/owners must stay alive until async metadata callbacks have - // drained because completions can release waiters through those gates. - prefetcher_stop(self->prefetcher); - metadata_store_async_destroy(self->store_meta_async); - self->store_meta_async = NULL; - prefetcher_destroy(self->prefetcher); - self->prefetcher = NULL; - prefetch_cache_destroy(self->chunk_layout_cache); - self->chunk_layout_cache = NULL; - prefetch_cache_destroy(self->shard_index_cache); - self->shard_index_cache = NULL; - prefetch_cache_destroy(self->array_meta_cache); - self->array_meta_cache = NULL; - wave_pool_destroy(&self->wave_pool, cuda_skip); - render_job_pool_destroy(&self->render_jobs, cuda_skip); +#ifdef DAMACY_HAS_CUDA +#include "executor/cuda_executor.h" +#endif - free(self->staging); - self->staging = NULL; - lookahead_destroy(&self->lookahead); - batch_pool_destroy(&self->batch_pool, cuda_skip); - - planner_destroy(self->planner); - self->planner = NULL; - store_destroy(self->store_gds); - self->store_gds = NULL; - store_destroy(self->store_host); - self->store_host = NULL; - gpu_budget_destroy(self->budget); - self->budget = NULL; -} - -struct resolved_wave_geometry -{ - const struct input_transfer_ops* input; - struct wave_pool_sizing sizing; - struct gpu_budget_breakdown predicted; - uint32_t max_chunks_per_wave; - uint32_t max_substreams_per_chunk; - uint8_t host_buffer_waves; - uint8_t want_gds; -}; - -enum wave_geometry_step -{ - WAVE_GEOMETRY_SIZING, - WAVE_GEOMETRY_PREDICT, -}; - -static enum damacy_status -resolve_wave_geometry(const struct damacy_config* cfg, - uint64_t resolver_budget, - uint64_t runtime_chunk_cap, - struct resolved_wave_geometry* out, - enum wave_geometry_step* failed_step) -{ - *out = (struct resolved_wave_geometry){ 0 }; - out->max_chunks_per_wave = resolve_max_chunks_per_wave(cfg); - out->max_substreams_per_chunk = resolve_max_substreams_per_chunk(cfg); - out->host_buffer_waves = resolve_host_buffer_waves(cfg); - out->want_gds = resolve_enable_gds(cfg); - out->input = - out->want_gds ? input_transfer_gds() : input_transfer_host_staging(); - - const struct input_transfer_resources min_input = - input_transfer_resources(out->input, out->host_buffer_waves, 0); - enum damacy_status s = - wave_pool_resolve_sizing(out->max_chunks_per_wave, - out->max_substreams_per_chunk, - min_input.device_staging_buffers, - resolver_budget, - runtime_chunk_cap, - cfg->samples_per_batch, - &out->sizing); - if (s != DAMACY_OK) { - if (failed_step) - *failed_step = WAVE_GEOMETRY_SIZING; - return s; - } - - const struct input_transfer_resources input_resources = - input_transfer_resources( - out->input, out->host_buffer_waves, out->sizing.input_staging_per_wave); - s = gpu_budget_predict(cfg, - &input_resources, - out->sizing.dev_decompressed_per_wave, - &out->predicted); - if (s != DAMACY_OK && failed_step) - *failed_step = WAVE_GEOMETRY_PREDICT; - return s; -} +#include +#include enum damacy_status -damacy_create(const struct damacy_config* cfg, struct damacy** out) +damacy_pipeline_create(struct damacy_planner* planner, + struct damacy_executor* executor, + const struct damacy_batch_spec* output, + const struct damacy_queue_limits* queues, + struct damacy** out) { - enum damacy_status s = DAMACY_INVAL; - struct damacy* self = NULL; - struct ctx_guard cg = { 0 }; - - CHECK_SILENT(InvalidArg, out); + if (!out) + return DAMACY_INVAL; *out = NULL; - - s = validate_config(cfg); - if (s != DAMACY_OK) - return s; - - s = DAMACY_OOM; - self = (struct damacy*)calloc(1, sizeof(*self)); - CHECK(Fail, self); - self->cfg = *cfg; - self->failed_status = DAMACY_OK; - self->page_alignment = (uint64_t)platform_page_alignment(); - // -1 sentinel set before any goto Fail: 0 from calloc is a valid CUdevice. - self->retained_primary_device = -1; - self->retained_primary = NULL; + if (!planner || !executor || !queues || !output || planner->active || + executor->active || !queues->prepared_batches || + queues->prepared_batches > 1024 || + queues->lookahead_samples < output->samples_per_batch || + queues->lookahead_samples > UINT32_MAX - 4) + return DAMACY_INVAL; + int64_t shape[DAMACY_MAX_RANK + 1], strides[DAMACY_MAX_RANK + 1]; + uint64_t bytes; + enum damacy_status status = batch_spec_layout(output, shape, strides, &bytes); + if (status != DAMACY_OK) + return status; + struct damacy* self = calloc(1, sizeof(*self)); + if (!self) + return DAMACY_OOM; + int planner_claimed = 0, executor_claimed = 0; + int planner_started = 0, executor_started = 0; + self->output = *output; + self->queues = *queues; + self->planner = planner; + self->executor = executor; stats_init(&self->stats); - - s = DAMACY_CUDA; - CU(Fail, cuInit(0)); - - CUcontext caller_ctx = NULL; - CU(Fail, cuCtxGetCurrent(&caller_ctx)); - if (cfg->device >= 0) { - if (caller_ctx) { - CUdevice cur_dev; - CU(Fail, cuCtxGetDevice(&cur_dev)); - if ((int)cur_dev != cfg->device) { - s = DAMACY_INVAL; - log_error("damacy_create: Config.device=%d but a CUcontext is " - "already current on device %d — likely a missing " - "cuCtxSetCurrent / torch.cuda.set_device(%d)", - cfg->device, - (int)cur_dev, - cfg->device); - goto Fail; - } - } - CUdevice dev; - CU(Fail, cuDeviceGet(&dev, cfg->device)); - CUcontext primary = NULL; - CU(Fail, cuDevicePrimaryCtxRetain(&primary, dev)); - self->retained_primary_device = cfg->device; - self->retained_primary = primary; - self->worker_ctx = primary; - self->cuda_device = cfg->device; - } else { - if (!caller_ctx) { - log_error("damacy_create: no CUcontext is current on calling thread"); - s = DAMACY_INVAL; - goto Fail; - } - CUdevice dev; - CU(Fail, cuCtxGetDevice(&dev)); - self->cuda_device = (int)dev; - self->worker_ctx = caller_ctx; - } - - s = ctx_guard_enter(self, &cg); - if (s != DAMACY_OK) + self->plans = calloc(queues->prepared_batches, sizeof(*self->plans)); + if (!self->plans) { + status = DAMACY_OOM; goto Fail; - - // Resolve the GPU's host-NUMA node now that we know the CUdevice. The - // resolved plan is consumed by the wave_pool_init scope below and by - // the store / scheduler worker threads at startup. - { - CUdevice cu_dev; - if (cuCtxGetDevice(&cu_dev) == CUDA_SUCCESS) { - numa_init( - cfg->tuning.numa_strategy, cfg->tuning.numa_node, cu_dev, &self->numa); - } else { - // Should never happen — ctx_guard_enter just pushed our ctx, or - // the caller's ctx is live. Be safe; treat as disabled. - self->numa.node = -1; - } } - - const uint64_t max_gpu = cfg->tuning.max_gpu_memory_bytes; - const uint64_t runtime_chunk_cap = resolve_max_chunk_uncompressed(cfg); - self->budget = gpu_budget_new(max_gpu); - if (!self->budget) { - s = DAMACY_OOM; + int expected = 0; + if (!atomic_compare_exchange_strong(&planner->active, &expected, 1)) { + status = DAMACY_INVAL; goto Fail; } - - // Carve out the double-buffered batch-output pool before sizing - // wave-resident buffers; the resolver is greedy, so without this - // reservation the lazy pool has no room at first push. - uint64_t pool_reserve = 0; - { - uint64_t pool_bytes = 0; - CHECK(Fail, - (s = resolve_sample_volume_bytes(cfg, &pool_bytes)) == DAMACY_OK); - pool_reserve = 2ull * pool_bytes; - } - if (pool_reserve >= max_gpu) { - log_error( - "damacy: batch-output pool reserve=%llu >= " - "max_gpu_memory_bytes=%llu; nothing left for wave-resident " - "buffers (sample_shape × samples_per_batch × dtype_bpe × 2 exceeds cap)", - (unsigned long long)pool_reserve, - (unsigned long long)max_gpu); - s = DAMACY_BUDGET; + planner_claimed = 1; + expected = 0; + if (!atomic_compare_exchange_strong(&executor->active, &expected, 1)) { + status = DAMACY_INVAL; goto Fail; } - const uint64_t resolver_budget = max_gpu - pool_reserve; - - struct resolved_wave_geometry geom = { 0 }; - s = - resolve_wave_geometry(cfg, resolver_budget, runtime_chunk_cap, &geom, NULL); - if (s != DAMACY_OK) + executor_claimed = 1; + status = planner->ops->start(planner, output, queues); + if (status != DAMACY_OK) goto Fail; - { - gpu_budget_commit(self->budget, geom.predicted.total); - log_debug("damacy: resolved geometry from max_gpu_memory_bytes=%llu " - "(pool_reserve=%llu, resolver_budget=%llu): " - "input_staging_per_wave=%llu dev_decompressed_per_wave=%llu " - "initial_nvcomp_temp=%llu predicted_total=%llu", - (unsigned long long)max_gpu, - (unsigned long long)pool_reserve, - (unsigned long long)resolver_budget, - (unsigned long long)geom.sizing.input_staging_per_wave, - (unsigned long long)geom.sizing.dev_decompressed_per_wave, - (unsigned long long)geom.predicted.nvcomp_temp, - (unsigned long long)geom.predicted.total); - // Resolver guarantees this fits; assert defensively in case the - // accounting drifts. A breach here is a bug, not user input. - if (gpu_budget_committed(self->budget) > gpu_budget_max(self->budget)) { - log_error( - "damacy: post-resolve drift: total=%llu cap=%llu " - "(dev_compressed=%llu dev_decompressed=%llu " - "blosc1_meta=%llu fanout_soa=%llu nvcomp_temp=%llu batch_meta=%llu)", - (unsigned long long)geom.predicted.total, - (unsigned long long)gpu_budget_max(self->budget), - (unsigned long long)geom.predicted.dev_compressed, - (unsigned long long)geom.predicted.dev_decompressed, - (unsigned long long)geom.predicted.blosc1_meta, - (unsigned long long)geom.predicted.fanout_soa, - (unsigned long long)geom.predicted.nvcomp_temp, - (unsigned long long)geom.predicted.batch_metadata); - s = DAMACY_BUDGET; - goto Fail; - } - } - - s = DAMACY_OOM; - - // Sample.uri is absolute; fs store joins root+key, so empty root is a - // pass-through. - { - struct store_fs_config sc = { - .root = "", - .nthreads = (int)cfg->tuning.n_io_threads, - .affinity = &self->numa, - // Worst case: every staging slot full of unfused single-chunk - // reads, plus slack for the stray reads that share the pool. - .max_inflight_reads = (uint32_t)geom.host_buffer_waves * - geom.max_chunks_per_wave + - DAMACY_READ_JOB_SLACK, - }; - self->store_host = store_fs_create(&sc); - CHECK(Fail, self->store_host); - } - self->store_meta_async = - metadata_store_async_create((int)resolve_metadata_io_concurrency(cfg), - &self->numa, - &cfg->debug.metadata_latency); - CHECK(Fail, self->store_meta_async); - if (geom.want_gds) { - struct store_fs_gds_config sc = { - .root = "", - .fd_cache_capacity = 0, - }; - self->store_gds = store_fs_gds_create(&sc); - if (!self->store_gds) { - s = DAMACY_INVAL; - goto Fail; - } - } - - array_meta_async_fetcher_init(&self->array_meta_async_fetcher, - self->store_meta_async); - { - struct prefetch_cache_config array_meta_cache_cfg = { - .capacity = cfg->tuning.n_array_meta_cache, - .max_probe = 16, - .knob_name = "n_array_meta_cache", - .ops = &array_meta_ops, - .async_fetcher = &self->array_meta_async_fetcher.base, - }; - self->array_meta_cache = prefetch_cache_create(&array_meta_cache_cfg); - CHECK(Fail, self->array_meta_cache); - } - shard_index_async_fetcher_init(&self->shard_index_async_fetcher, - self->store_meta_async, - self->array_meta_cache); - { - struct prefetch_cache_config shard_index_cache_cfg = { - .capacity = cfg->tuning.n_shard_index_cache, - .max_probe = 16, - .knob_name = "n_shard_index_cache", - .ops = &shard_index_ops, - .async_fetcher = &self->shard_index_async_fetcher.base, - }; - self->shard_index_cache = prefetch_cache_create(&shard_index_cache_cfg); - CHECK(Fail, self->shard_index_cache); - } - chunk_layout_async_fetcher_init(&self->chunk_layout_async_fetcher, - self->store_meta_async, - self->array_meta_cache, - self->shard_index_cache, - geom.max_substreams_per_chunk); - { - struct prefetch_cache_config chunk_layout_cache_cfg = { - .capacity = cfg->tuning.n_chunk_layout_cache, - .max_probe = 16, - .knob_name = "n_chunk_layout_cache", - .ops = &chunk_layout_ops, - .async_fetcher = &self->chunk_layout_async_fetcher.base, - }; - self->chunk_layout_cache = prefetch_cache_create(&chunk_layout_cache_cfg); - CHECK(Fail, self->chunk_layout_cache); - } - - for (int b = 0; b < DAMACY_N_BATCH_SLOTS; ++b) - CHECK(Fail, - batch_slot_init(&self->batch_pool.slots[b], cfg->samples_per_batch) == - 0); - for (int b = 0; b < DAMACY_N_BATCH_SLOTS; ++b) { - struct render_job* job = - render_job_pool_for_batch_slot(&self->render_jobs, (uint16_t)b); - CHECK(Fail, render_job_init(job, cfg->samples_per_batch) == 0); - } - - s = DAMACY_OOM; - // Keep first-touch allocations on the GPU's NUMA node. - { - struct platform_cpu_mask saved_affinity; - numa_scope_enter(&self->numa, &saved_affinity); - int wp_rc = - wave_pool_init(&self->wave_pool, - &self->batch_pool, - &self->render_jobs, - geom.want_gds ? self->store_gds : self->store_host, - &self->stats, - cfg->dtype, - geom.host_buffer_waves, - geom.max_chunks_per_wave, - geom.max_substreams_per_chunk, - geom.sizing.input_staging_per_wave, - geom.sizing.dev_decompressed_per_wave, - runtime_chunk_cap, - geom.input, - cfg->debug.bypass_decode, - self->budget); - numa_scope_exit(&saved_affinity); - CHECK(Fail, wp_rc == 0); - } - if (geom.want_gds) - log_info("damacy: input transfer via cuFile / GDS"); - - struct planner_config pcfg = { - .array_meta_cache = self->array_meta_cache, - .chunk_layout_cache = self->chunk_layout_cache, - .shard_index_cache = self->shard_index_cache, - .dst_dtype = cfg->dtype, - .page_alignment = self->page_alignment, - .max_chunk_uncompressed_bytes = runtime_chunk_cap, - .read_op_max_bytes = resolve_max_read_op_bytes(cfg), - .max_chunks_per_wave = geom.max_chunks_per_wave, - .max_substreams_per_chunk = geom.max_substreams_per_chunk, - }; - CHECK(Fail, planner_create(&pcfg, &self->planner) == DAMACY_OK); - - CHECK(Fail, lookahead_init(&self->lookahead, cfg->lookahead_samples) == 0); - - { - struct prefetcher_config prefetcher_cfg = { - .lookahead = &self->lookahead, - .array_meta_cache = self->array_meta_cache, - .shard_index_cache = self->shard_index_cache, - .chunk_layout_cache = self->chunk_layout_cache, - .capacity = cfg->lookahead_samples, - .owner_capacity = cfg->lookahead_samples + 4u, - .max_shards_per_sample = cfg->tuning.max_shards_per_sample, - }; - self->prefetcher = prefetcher_create(&prefetcher_cfg); - CHECK(Fail, self->prefetcher); - } - - self->staging = (struct prefetcher_ready*)calloc( - 2u * (size_t)cfg->samples_per_batch, sizeof(struct prefetcher_ready)); - CHECK(Fail, self->staging); - - self->handle.d = self; - - // Start the prefetcher worker before the scheduler so the scheduler's - // first tick can see ready batches. - CHECK(Fail, prefetcher_start(self->prefetcher) == 0); - - // Spawn the worker last — everything it touches must already exist. - self->sched = scheduler_create( - damacy_scheduler_step, self, DAMACY_POP_POLL_NS, &self->numa); + planner_started = 1; + status = executor->ops->start(executor, output, &self->stats); + if (status != DAMACY_OK) + goto Fail; + executor_started = 1; + self->device = + executor->device_type == DAMACY_DEVICE_CUDA ? executor->device_id : -1; + self->sched = + scheduler_create(damacy_scheduler_step, + self, + 10000, + NULL, + &(struct scheduler_hooks){ damacy_scheduler_enter, + damacy_scheduler_leave }); if (!self->sched) { - s = DAMACY_OOM; + status = DAMACY_OOM; goto Fail; } - - ctx_guard_exit(&cg); *out = self; return DAMACY_OK; - Fail: - if (self) { - // destroy_inner under the pushed ctx, then pop, then release primary. - destroy_inner(self, 0); - ctx_guard_exit(&cg); - if (self->retained_primary_device >= 0) - cuDevicePrimaryCtxRelease((CUdevice)self->retained_primary_device); - free(self); - } - return s; + if (executor_started) + executor->ops->stop(executor); + if (planner_started) + planner->ops->stop(planner); + if (executor_claimed) + executor->active = 0; + if (planner_claimed) + planner->active = 0; + free(self->plans); + free(self); + return status; +} -InvalidArg: - return DAMACY_INVAL; +void +damacy_shutdown(struct damacy* self) +{ + if (!self || self->stopped) + return; + scheduler_lock(self->sched); + if (self->stopping) { + while (!self->stopped) + scheduler_wait(self->sched); + scheduler_unlock(self->sched); + return; + } + self->planner->ops->stats(self->planner, &self->stats); + self->executor->ops->stats(self->executor, &self->stats); + self->stopping = 1; + scheduler_broadcast(self->sched); + while (self->pop_calls) + scheduler_wait(self->sched); + scheduler_unlock(self->sched); + scheduler_stop(self->sched); + self->planner->ops->stop(self->planner); + self->executor->ops->stop(self->executor); + self->planner->active = self->executor->active = 0; + for (uint32_t i = 0; i < self->queues.prepared_batches; ++i) { + prepared_plan_destroy(self->plans[i]); + self->plans[i] = NULL; + } + scheduler_lock(self->sched); + self->plan_count = 0; + self->stopped = 1; + scheduler_broadcast(self->sched); + scheduler_unlock(self->sched); } void @@ -485,118 +134,133 @@ damacy_destroy(struct damacy* self) { if (!self) return; - - struct ctx_guard cg = { 0 }; - enum damacy_status gs = ctx_guard_enter(self, &cg); - if (gs != DAMACY_OK) { - // Ctx no longer pushable (device reset, primary released elsewhere): - // leak CUDA state and walk the teardown list with skip=1. - log_warn("damacy_destroy: ctx_guard_enter failed (status=%d); " - "leaking CUDA resources", - (int)gs); - destroy_inner(self, 1); - } else { - destroy_inner(self, 0); - ctx_guard_exit(&cg); + damacy_shutdown(self); + scheduler_destroy(self->sched); + if (self->owns_components) { + damacy_executor_destroy(self->executor); + damacy_planner_destroy(self->planner); + damacy_metadata_destroy(self->owned_metadata); + damacy_metadata_reader_destroy(self->owned_metadata_reader); + damacy_reader_destroy(self->owned_reader); } - if (self->retained_primary_device >= 0) - cuDevicePrimaryCtxRelease((CUdevice)self->retained_primary_device); + free(self->plans); free(self); } int -damacy_get_device(const struct damacy* d) +damacy_get_device(const struct damacy* self) { - return d ? d->cuda_device : -1; + return self ? self->device : -1; } -// Test-only hook. Overwrites gpu_bytes_committed so unit tests can drive -// the observe-and-grow OOM path without having to fabricate a workload -// that escapes the resolver's worst-case reservation. Returns the prior -// value. Not declared in damacy.h — tests forward-declare it with -// extern; production code must not call it. -uint64_t -damacy_set_gpu_bytes_committed_for_test(struct damacy* self, uint64_t v) +enum damacy_status +damacy_create(const struct damacy_config* config, struct damacy** out) { - if (!self) - return 0; - return gpu_budget_set_committed_for_test(self->budget, v); + if (!out) + return DAMACY_INVAL; + *out = NULL; + enum damacy_status status = validate_config(config); + if (status != DAMACY_OK) + return status; + struct damacy_reader* reader = NULL; + struct damacy_metadata_reader* metadata_reader = NULL; + struct damacy_metadata* metadata = NULL; + struct damacy_planner* planner = NULL; + struct damacy_executor* executor = NULL; + const struct damacy_tuning* tuning = &config->tuning; +#ifdef DAMACY_HAS_CUDA + struct numa_resolved affinity; + struct platform_cpu_mask saved; + cuda_resolve_numa(config, &affinity); + numa_scope_enter(&affinity, &saved); +#endif + status = damacy_file_reader_create(tuning->n_io_threads, + (uint32_t)tuning->host_buffer_waves * + tuning->max_chunks_per_wave + + DAMACY_READ_JOB_SLACK, + &reader); + if (status != DAMACY_OK) + goto Done; + status = damacy_file_metadata_reader_create(tuning->metadata_io_concurrency, + &config->debug.metadata_latency, + &metadata_reader); + if (status != DAMACY_OK) + goto Done; + status = damacy_zarr_metadata_create( + metadata_reader, + &(struct damacy_metadata_cache_config){ + .array_entries = tuning->n_array_meta_cache, + .shard_entries = tuning->n_shard_index_cache }, + &metadata); + if (status != DAMACY_OK) + goto Done; + status = damacy_chunk_planner_create( + metadata, + &(struct damacy_plan_limits){ + .max_chunks = DAMACY_MAX_CHUNKS_PER_BATCH, + .max_chunk_bytes = tuning->max_chunk_uncompressed_bytes, + .max_shards_per_sample = tuning->max_shards_per_sample, + .max_plan_bytes = 64ull << 20 }, + &planner); + if (status != DAMACY_OK) + goto Done; + status = damacy_cuda_executor_create( + reader, + &(struct damacy_cuda_config){ + .device = config->device, + .max_gpu_memory_bytes = tuning->max_gpu_memory_bytes, + .max_chunk_bytes = tuning->max_chunk_uncompressed_bytes, + .max_read_bytes = tuning->max_read_op_bytes, + .max_chunks_per_wave = tuning->max_chunks_per_wave, + .max_substreams_per_chunk = tuning->max_substreams_per_chunk, + .host_buffer_waves = tuning->host_buffer_waves, + .chunk_layout_entries = tuning->n_chunk_layout_cache, + .numa_strategy = tuning->numa_strategy, + .numa_node = tuning->numa_node, + .enable_gds = tuning->enable_gds, + .bypass_decode = config->debug.bypass_decode }, + &executor); + if (status != DAMACY_OK) + goto Done; + struct damacy_batch_spec output = { .dtype = config->dtype, + .sample_rank = config->sample_rank, + .samples_per_batch = + config->samples_per_batch }; + memcpy( + output.sample_shape, config->sample_shape, sizeof(output.sample_shape)); + status = damacy_pipeline_create( + planner, + executor, + &output, + &(struct damacy_queue_limits){ + .lookahead_samples = config->lookahead_samples, .prepared_batches = 2 }, + out); + if (status == DAMACY_OK) { + (*out)->owns_components = 1; + (*out)->owned_reader = reader; + (*out)->owned_metadata_reader = metadata_reader; + (*out)->owned_metadata = metadata; +#ifdef DAMACY_HAS_CUDA + numa_scope_exit(&saved); +#endif + return DAMACY_OK; + } +Done: +#ifdef DAMACY_HAS_CUDA + numa_scope_exit(&saved); +#endif + damacy_executor_destroy(executor); + damacy_planner_destroy(planner); + damacy_metadata_destroy(metadata); + damacy_metadata_reader_destroy(metadata_reader); + damacy_reader_destroy(reader); + return status; } -void -damacy_config_describe(const struct damacy_config* cfg) +#ifdef DAMACY_HAS_CUDA +uint64_t +damacy_set_gpu_bytes_committed_for_test(struct damacy* self, uint64_t value) { - if (!cfg) { - log_info("damacy_config_describe: NULL config"); - return; - } - const uint64_t max_gpu = cfg->tuning.max_gpu_memory_bytes; - const uint64_t runtime_chunk_cap = resolve_max_chunk_uncompressed(cfg); - uint64_t pool_reserve = 0; - { - uint64_t pool_bytes = 0; - // resolve_sample_volume_bytes rejects rank=0 / non-positive dims; on - // those, leave pool_reserve at 0 so describe still prints useful - // info for the rest of the geometry. - enum damacy_status pvs = resolve_sample_volume_bytes(cfg, &pool_bytes); - if (pvs == DAMACY_OK) - pool_reserve = 2ull * pool_bytes; - else - log_info( - "damacy_config_describe: resolve_sample_volume_bytes failed (%s); " - "pool_reserve=0", - damacy_status_str(pvs)); - } - const uint64_t resolver_budget = - pool_reserve < max_gpu ? max_gpu - pool_reserve : 0; - log_info("damacy_config_describe: max_gpu_memory_bytes=%llu " - "(pool_reserve=%llu, resolver_budget=%llu, " - "max_chunk_uncompressed_bytes=%llu, samples_per_batch=%u)", - (unsigned long long)max_gpu, - (unsigned long long)pool_reserve, - (unsigned long long)resolver_budget, - (unsigned long long)runtime_chunk_cap, - (unsigned)cfg->samples_per_batch); - - struct resolved_wave_geometry geom = { 0 }; - enum wave_geometry_step failed_step = WAVE_GEOMETRY_SIZING; - enum damacy_status rs = resolve_wave_geometry( - cfg, resolver_budget, runtime_chunk_cap, &geom, &failed_step); - if (rs != DAMACY_OK) { - const char* where = failed_step == WAVE_GEOMETRY_PREDICT - ? "gpu_budget_predict" - : "wave_pool_resolve_sizing"; - log_info( - "damacy_config_describe: %s failed (%s)", where, damacy_status_str(rs)); - return; - } - log_info("damacy_config_describe: input_staging_per_wave=%llu " - "dev_decompressed_per_wave=%llu", - (unsigned long long)geom.sizing.input_staging_per_wave, - (unsigned long long)geom.sizing.dev_decompressed_per_wave); - log_info("damacy_config_describe: dev_compressed=%llu dev_decompressed=%llu " - "blosc1_meta=%llu fanout_soa=%llu " - "nvcomp_temp=%llu batch_metadata=%llu", - (unsigned long long)geom.predicted.dev_compressed, - (unsigned long long)geom.predicted.dev_decompressed, - (unsigned long long)geom.predicted.blosc1_meta, - (unsigned long long)geom.predicted.fanout_soa, - (unsigned long long)geom.predicted.nvcomp_temp, - (unsigned long long)geom.predicted.batch_metadata); - // predicted.total is the *initial* allocation (initial fanout / decoder - // floors); sizing.worst_case_total_bytes is the post-grow worst case - // the resolver pre-reserved against the cap. Their difference is the - // grow-time headroom; the cap minus the worst case is unused slack. - const uint64_t initial_alloc = geom.predicted.total; - const uint64_t worst_case = geom.sizing.worst_case_total_bytes; - const uint64_t reserved_for_grow = - worst_case > initial_alloc ? worst_case - initial_alloc : 0; - const uint64_t slack = max_gpu > worst_case ? max_gpu - worst_case : 0; - log_info("damacy_config_describe: initial_alloc=%llu reserved_for_grow=%llu " - "worst_case_total=%llu slack=%llu cap=%llu", - (unsigned long long)initial_alloc, - (unsigned long long)reserved_for_grow, - (unsigned long long)worst_case, - (unsigned long long)slack, - (unsigned long long)max_gpu); + return self ? cuda_executor_set_budget(self->executor, value) : 0; } +#endif diff --git a/src/damacy_pipeline.h b/src/damacy_pipeline.h new file mode 100644 index 00000000..2aed9149 --- /dev/null +++ b/src/damacy_pipeline.h @@ -0,0 +1,112 @@ +#pragma once + +#include "damacy.h" + +#ifdef __cplusplus +extern "C" +{ +#endif + + struct damacy_reader; + struct damacy_metadata_reader; + struct damacy_metadata; + struct damacy_planner; + struct damacy_executor; + + struct damacy_batch_spec + { + enum damacy_dtype dtype; + int64_t sample_shape[DAMACY_MAX_RANK]; + uint8_t sample_rank; + uint32_t samples_per_batch; + }; + + struct damacy_queue_limits + { + uint32_t lookahead_samples; + uint32_t prepared_batches; + }; + + struct damacy_plan_limits + { + uint32_t max_chunks; + uint32_t max_chunk_bytes; + uint32_t max_shards_per_sample; + uint64_t max_plan_bytes; + }; + + struct damacy_metadata_cache_config + { + uint32_t array_entries; + uint32_t shard_entries; + }; + + struct damacy_cpu_config + { + uint32_t decode_workers; + uint32_t max_encoded_chunk_bytes; + uint32_t max_decoded_chunk_bytes; + uint64_t max_memory_bytes; + }; + + struct damacy_cuda_config + { + int device; + uint64_t max_gpu_memory_bytes; + uint32_t max_chunk_bytes; + uint64_t max_read_bytes; + uint32_t max_chunks_per_wave; + uint32_t max_substreams_per_chunk; + uint8_t host_buffer_waves; + uint32_t chunk_layout_entries; + enum damacy_numa_strategy numa_strategy; + int numa_node; + enum damacy_gds_mode enable_gds; + uint8_t bypass_decode; + }; + + enum damacy_status damacy_file_reader_create(uint32_t workers, + uint32_t max_inflight_reads, + struct damacy_reader** out); + void damacy_reader_destroy(struct damacy_reader* reader); + + enum damacy_status damacy_file_metadata_reader_create( + uint32_t concurrency, + const struct damacy_latency_model* latency, + struct damacy_metadata_reader** out); + void damacy_metadata_reader_destroy(struct damacy_metadata_reader* reader); + + enum damacy_status damacy_zarr_metadata_create( + struct damacy_metadata_reader* reader, + const struct damacy_metadata_cache_config* cache, + struct damacy_metadata** out); + void damacy_metadata_destroy(struct damacy_metadata* metadata); + + enum damacy_status damacy_chunk_planner_create( + struct damacy_metadata* metadata, + const struct damacy_plan_limits* limits, + struct damacy_planner** out); + void damacy_planner_destroy(struct damacy_planner* planner); + + enum damacy_status damacy_cpu_executor_create( + struct damacy_reader* reader, + const struct damacy_cpu_config* config, + struct damacy_executor** out); + enum damacy_status damacy_cuda_executor_create( + struct damacy_reader* reader, + const struct damacy_cuda_config* config, + struct damacy_executor** out); + void damacy_executor_destroy(struct damacy_executor* executor); + + // Components are borrowed until shutdown; each planner/executor may serve + // one active pipeline. Readers and metadata outlive their dependents. + enum damacy_status damacy_pipeline_create( + struct damacy_planner* planner, + struct damacy_executor* executor, + const struct damacy_batch_spec* output, + const struct damacy_queue_limits* queues, + struct damacy** out); + +#ifdef __cplusplus +} +#endif diff --git a/src/damacy_plan.c b/src/damacy_plan.c index 4483f6a0..1a592d2e 100644 --- a/src/damacy_plan.c +++ b/src/damacy_plan.c @@ -1,333 +1,28 @@ -#include "damacy.h" - -#include "damacy_config.h" #include "damacy_internal.h" -#include "damacy_stats.h" -#include "platform/platform.h" -#include "util/cuda_check.h" -#include "util/prelude.h" - -#include -#include - -// Lazy batch-output pool sizing + GPU-budget enforcement. Geometry is -// fixed by cfg->sample_shape at create-time; this only allocates the -// device buffers on first push. Idempotent. -static enum damacy_status -batch_pool_allocate(struct damacy* self) -{ - struct damacy_batch_pool* pool = &self->batch_pool; - if (pool->allocated) - return DAMACY_OK; - - enum damacy_status s = - batch_pool_compute_layout(pool, - self->cfg.sample_shape, - self->cfg.sample_rank, - self->cfg.samples_per_batch, - damacy_dtype_bpe(self->cfg.dtype)); - if (s != DAMACY_OK) - return s; - - const uint64_t need = 2ull * pool->n_bytes; - s = gpu_budget_try_commit(self->budget, need, "batch-output pool"); - if (s != DAMACY_OK) - return s; - - s = batch_pool_alloc_dev(pool); - if (s != DAMACY_OK) { - gpu_budget_release(self->budget, need); - return s; - } - return DAMACY_OK; -} - -// --- plan: accumulate [locked] → run sealed [unlocked] → commit [locked] ---- -// Ready samples are staged cheaply while the batch is still open. Only a -// sealed full batch enters BATCH_PLANNING; plan_run then does planner -// CPU work + sample_plans upload off scheduler_lock. Pop treats -// BATCH_PLANNING as "planner work is outstanding". - -static void -free_stage_sample(struct planner_sample* s) -{ - if (!s) - return; - free(s->h_shards); - free((char*)s->uri); - *s = (struct planner_sample){ 0 }; -} - -static void -free_slot_stage_samples(struct damacy_batch_slot* slot) -{ - if (!slot || !slot->stage_samples) - return; - for (uint32_t i = 0; i < slot->n_samples; ++i) - free_stage_sample(&slot->stage_samples[i]); -} -static enum damacy_status -fail_plan_commit_slot(struct damacy* self, - struct damacy_batch_slot* slot, - struct render_job* job, - enum damacy_status status, - int* changed) -{ - free_slot_stage_samples(slot); - batch_slot_reset_for_reuse(slot); - render_job_reset(job); - self->failed_status = status; - if (changed) - *changed = 1; - return status; -} +#include "damacy_stats.h" enum damacy_status -plan_reserve(struct damacy* self, - uint16_t slot_idx, - struct prefetcher_ready* ready, - uint32_t n_ready, - int close_batch) -{ - if (n_ready == 0) - return DAMACY_OK; - struct damacy_batch_slot* slot = &self->batch_pool.slots[slot_idx]; - if (slot->state != BATCH_FREE && slot->state != BATCH_ACCUMULATING) - return DAMACY_INVAL; - uint32_t begin = slot->n_samples; - if (slot->state == BATCH_FREE) { - if (begin != 0) - return DAMACY_INVAL; - slot->batch_id = self->next_batch_id; - slot->sample_seq_begin = ready[0].sample_seq; - } else if (slot->batch_id != self->next_batch_id) { - return DAMACY_INVAL; - } - if (begin + n_ready > self->cfg.samples_per_batch) - return DAMACY_INVAL; - - for (uint32_t i = 0; i < n_ready; ++i) { - // advance_from_shard absorbs NOTFOUND at the prefetcher level; a - // slot-level ERROR here is always a real error. - if (ready[i].result == PREFETCHER_RESULT_ERROR) { - enum damacy_status es = ready[i].err_code - ? (enum damacy_status)ready[i].err_code - : DAMACY_INVAL; - self->failed_status = es; - return es; - } - if (ready[i].sample_seq != slot->sample_seq_begin + begin + i) - return DAMACY_INVAL; - } - - if (close_batch) { - enum damacy_status status = batch_pool_allocate(self); - if (status != DAMACY_OK) { - self->failed_status = status; +pipeline_prepare(struct damacy* self, int* changed) +{ + while (self->plan_count < self->queues.prepared_batches) { + struct prepared_plan* plan = NULL; + struct platform_clock clock = { 0 }; + platform_toc(&clock); + enum damacy_status status = self->planner->ops->next(self->planner, &plan); + if (status == DAMACY_AGAIN) + return DAMACY_OK; + if (status != DAMACY_OK) return status; - } - } - for (uint32_t i = 0; i < n_ready; ++i) { - struct planner_sample* dst = &slot->stage_samples[begin + i]; - dst->uri = ready[i].uri; - dst->aabb = ready[i].aabb; - dst->h_meta = ready[i].h_meta; - dst->h_shards = ready[i].h_shards; - dst->n_shards = ready[i].n_shards; - dst->h_layout = ready[i].h_layout; - ready[i].uri = NULL; - ready[i].h_shards = NULL; - ready[i].n_shards = 0; - } - slot->n_samples = begin + n_ready; - slot->planning_close_batch = close_batch; - slot->state = close_batch ? BATCH_PLANNING : BATCH_ACCUMULATING; - return DAMACY_OK; -} - -// Returns elapsed ms via *out_elapsed_ms so the metric can be recorded -// in plan_commit, which runs under scheduler_lock (stats are read by -// damacy_stats_get under the same lock). -enum damacy_status -plan_run(struct damacy* self, uint16_t slot_idx, float* out_elapsed_ms) -{ - struct render_job* job = - render_job_pool_for_batch_slot(&self->render_jobs, slot_idx); - CHECK(InvalidArg, job); - struct damacy_batch_slot* slot = &self->batch_pool.slots[slot_idx]; - CHECK(InvalidArg, slot->state == BATCH_PLANNING); - render_job_reset(job); - struct planner_output plan_out = - render_job_planner_output(job, self->cfg.samples_per_batch); - struct platform_clock plan_clock = { 0 }; - platform_toc(&plan_clock); - enum damacy_status status = planner_plan(self->planner, - slot->stage_samples, - slot->n_samples, - slot_idx, - self->batch_pool.strides, - self->batch_pool.rank, - &plan_out); - *out_elapsed_ms = platform_toc(&plan_clock) * 1000.0f; - if (status != DAMACY_OK) - return status; - render_job_commit_plan(job, slot_idx, slot->batch_id, &plan_out); - return render_job_upload_sample_plans(job); -InvalidArg: - return DAMACY_INVAL; -} - -// *changed (nullable): OR-set on every BATCH transition. -enum damacy_status -plan_commit(struct damacy* self, - uint16_t slot_idx, - enum damacy_status run_status, - float elapsed_ms, - int* changed) -{ - metric_record(&self->stats.plan, elapsed_ms, 0, 0); - struct render_job* job = - render_job_pool_for_batch_slot(&self->render_jobs, slot_idx); - if (!job) { - self->failed_status = DAMACY_INVAL; - return DAMACY_INVAL; - } - struct damacy_batch_slot* slot = &self->batch_pool.slots[slot_idx]; - if (run_status != DAMACY_OK) - return fail_plan_commit_slot(self, slot, job, run_status, changed); - if (!slot->planning_close_batch) - return fail_plan_commit_slot(self, slot, job, DAMACY_INVAL, changed); - slot->n_chunks = job->n_chunks; - slot->chunks_remaining = (int32_t)slot->n_chunks; - - free_slot_stage_samples(slot); - prefetcher_advance_watermark(self->prefetcher, - slot->sample_seq_begin + slot->n_samples); - self->next_batch_id++; - slot->planning_close_batch = 0; - slot->state = BATCH_RENDERING; - self->stats.chunks_planned += slot->n_chunks; - self->stats.chunks_to_load += job->n_chunks_to_load; - self->stats.reads_issued += job->n_loads_issued; - if (changed) + metric_record(&self->stats.plan, platform_toc(&clock) * 1000, 0, 0); + uint32_t index = + (self->plan_head + self->plan_count) % self->queues.prepared_batches; + self->plans[index] = plan; + ++self->plan_count; + self->stats.chunks_planned += plan->n_uses; + for (uint32_t i = 0; i < plan->n_uses; ++i) + self->stats.chunks_to_load += !plan->chunks[plan->uses[i].chunk].missing; *changed = 1; - - if (slot->n_chunks == 0) { - // Degenerate batch: zero the output and skip to READY. - if (slot->deferred_reuse_pending) { - cuStreamSynchronize(self->wave_pool.stream_post); - slot->deferred_reuse_pending = 0; - } - if (cuMemsetD8(CUDPTR(slot->dev_ptr), 0, self->batch_pool.n_bytes) != - CUDA_SUCCESS) - return fail_plan_commit_slot(self, slot, job, DAMACY_CUDA, changed); - slot->state = BATCH_READY; - render_job_finish(job); } - // Non-degenerate path doesn't clear the flag: assemble on FIFO stream_post - // already inherits the cuStreamWaitEvent, so host-side consumption is moot. return DAMACY_OK; } - -static enum damacy_status -run_sealed_plan(struct damacy* self, uint16_t slot_idx, int* changed) -{ - scheduler_unlock(self->sched); - float plan_ms = 0.f; - enum damacy_status rs = plan_run(self, slot_idx, &plan_ms); - scheduler_lock(self->sched); - return plan_commit(self, slot_idx, rs, plan_ms, changed); -} - -static uint32_t -planning_capacity_locked(struct damacy* self) -{ - if (any_batch_planning(&self->batch_pool)) - return 0; - uint32_t cap = 0; - int open_slot = find_accumulating_batch_slot(&self->batch_pool); - if (open_slot >= 0) { - struct damacy_batch_slot* slot = &self->batch_pool.slots[open_slot]; - if (slot->n_samples < self->cfg.samples_per_batch) - cap += self->cfg.samples_per_batch - slot->n_samples; - } - for (int s = 0; s < DAMACY_N_BATCH_SLOTS; ++s) { - if (self->batch_pool.slots[s].state == BATCH_FREE) - cap += self->cfg.samples_per_batch; - } - uint32_t staging_cap = 2u * self->cfg.samples_per_batch; - return cap < staging_cap ? cap : staging_cap; -} - -static int -next_plan_slot_locked(struct damacy* self) -{ - int open_slot = find_accumulating_batch_slot(&self->batch_pool); - if (open_slot >= 0) - return open_slot; - return find_free_batch_slot(&self->batch_pool); -} - -static void -free_ready_range(struct prefetcher_ready* ready, uint32_t n) -{ - for (uint32_t i = 0; i < n; ++i) - prefetcher_ready_free(&ready[i]); -} - -enum damacy_status -plan_ready_prefetch(struct damacy* self, int* changed) -{ - enum damacy_status status = DAMACY_OK; - for (;;) { - uint32_t cap = planning_capacity_locked(self); - if (cap == 0) - return DAMACY_OK; - - struct prefetcher_wave_ticket ticket = { 0 }; - if (!prefetcher_take_ready_wave( - self->prefetcher, cap, &ticket, self->staging)) { - return DAMACY_OK; - } - - uint32_t cursor = 0; - while (cursor < ticket.n_samples) { - int slot_idx = next_plan_slot_locked(self); - if (slot_idx < 0) { - free_ready_range(&self->staging[cursor], ticket.n_samples - cursor); - return DAMACY_INVAL; - } - - struct damacy_batch_slot* slot = &self->batch_pool.slots[slot_idx]; - uint32_t begin = slot->n_samples; - uint32_t room = self->cfg.samples_per_batch - begin; - uint32_t remaining = ticket.n_samples - cursor; - uint32_t n = remaining < room ? remaining : room; - int close_batch = begin + n == self->cfg.samples_per_batch; - - status = plan_reserve( - self, (uint16_t)slot_idx, &self->staging[cursor], n, close_batch); - if (status != DAMACY_OK) { - free_ready_range(&self->staging[cursor], ticket.n_samples - cursor); - return status; - } - free_ready_range(&self->staging[cursor], n); - if (changed) - *changed = 1; - - if (close_batch) { - status = run_sealed_plan(self, (uint16_t)slot_idx, changed); - if (status != DAMACY_OK) { - free_ready_range(&self->staging[cursor + n], - ticket.n_samples - cursor - n); - return status; - } - } - cursor += n; - } - - if (ticket.n_samples < cap) - return DAMACY_OK; - } -} diff --git a/src/damacy_pop.c b/src/damacy_pop.c index 75e1a36b..e8f06dd5 100644 --- a/src/damacy_pop.c +++ b/src/damacy_pop.c @@ -1,194 +1,81 @@ -#include "damacy.h" - #include "damacy_internal.h" + #include "damacy_stats.h" -#include "log/log.h" -#include "nvtx/nvtx.h" -#include "platform/platform.h" -#include "util/prelude.h" -#include #include -static void -release_slot_now(struct damacy_batch_slot* slot, struct render_job* job) -{ - batch_slot_reset_for_reuse(slot); - render_job_reset(job); -} - -static void -release_slot_after_event(struct damacy_batch_slot* slot, struct render_job* job) -{ - release_slot_now(slot, job); - slot->deferred_reuse_pending = 1; -} - enum damacy_status damacy_pop(struct damacy* self, struct damacy_batch** out) { - CHECK_SILENT(InvalidArg, self); - CHECK_SILENT(InvalidArg, out); + if (!out) + return DAMACY_INVAL; *out = NULL; - - // No ctx_guard: pop only touches batch-slot state. CUDA stays on the worker. - damacy_nvtx_range_push("damacy_pop"); - enum damacy_status r; + if (!self) + return DAMACY_INVAL; + if (self->stopped) + return DAMACY_SHUTDOWN; + enum damacy_status status; scheduler_lock(self->sched); + ++self->pop_calls; for (;;) { - if (self->failed_status != DAMACY_OK) { - r = self->failed_status; - goto Done; + if (self->stopping) { + status = DAMACY_SHUTDOWN; + break; } - int slot_idx = find_oldest_ready_slot(&self->batch_pool); - if (slot_idx >= 0) { - struct damacy_batch_slot* slot = &self->batch_pool.slots[slot_idx]; - slot->state = BATCH_HELD; - self->handle.slot_idx = (uint16_t)slot_idx; - self->handle.batch_id = slot->batch_id; - self->stats.batches_emitted++; - *out = &self->handle; - r = DAMACY_OK; - goto Done; + if (self->failed_status != DAMACY_OK) { + status = self->failed_status; + break; } - if (!any_wave_in_flight(&self->wave_pool) && - !any_slot_in_flight(&self->wave_pool) && - !any_batch_in_flight(&self->batch_pool) && - lookahead_size(&self->lookahead) == 0 && - prefetcher_in_flight(self->prefetcher) == 0 && - !prefetcher_has_ready(self->prefetcher)) { - r = DAMACY_AGAIN; - goto Done; + status = self->executor->ops->take(self->executor, out); + if (status != DAMACY_AGAIN) { + if (status == DAMACY_OK) { + (*out)->owner = self; + ++self->stats.batches_emitted; + } + break; } - struct platform_clock wait_clock = { 0 }; - platform_toc(&wait_clock); + if (!self->plan_count && !self->planner->ops->pending(self->planner) && + !self->executor->ops->busy(self->executor)) + break; + struct platform_clock clock = { 0 }; + platform_toc(&clock); SCHEDULER_WAIT_DIAG(self->sched, 5000); - metric_record( - &self->stats.pop_wait, platform_toc(&wait_clock) * 1000.0f, 0, 0); + metric_record(&self->stats.pop_wait, platform_toc(&clock) * 1000, 0, 0); } - -Done: + --self->pop_calls; + scheduler_broadcast(self->sched); scheduler_unlock(self->sched); - damacy_nvtx_range_pop(); - return r; - -InvalidArg: - return DAMACY_INVAL; + return status; } void -damacy_release(struct damacy* self, struct damacy_batch* b) +damacy_release(struct damacy* self, struct damacy_batch* batch) { - if (!self || !b) - return; - if (b != &self->handle) { - log_warn("damacy_release: foreign handle (not the active batch)"); - return; - } - uint16_t s = b->slot_idx; - if (s >= DAMACY_N_BATCH_SLOTS) { - log_warn("damacy_release: slot_idx=%u out of range", (unsigned)s); - return; - } - struct render_job* job = - render_job_pool_for_batch_slot(&self->render_jobs, s); - if (!job) - return; - scheduler_lock(self->sched); - if (self->batch_pool.slots[s].state != BATCH_HELD) { - log_warn("damacy_release: slot %u not HELD (state=%d); double release?", - (unsigned)s, - (int)self->batch_pool.slots[s].state); - scheduler_unlock(self->sched); - return; - } - release_slot_now(&self->batch_pool.slots[s], job); - scheduler_unlock(self->sched); + if (batch && batch->owner == self) + damacy_batch_release(batch); } enum damacy_status -damacy_release_event(struct damacy* self, struct damacy_batch* b, void* event) +damacy_release_event(struct damacy* self, + struct damacy_batch* batch, + void* event) { - // NULL event → degenerate to the immediate-release path. + if (!self || !batch || batch->owner != self) + return DAMACY_INVAL; if (!event) { - damacy_release(self, b); + damacy_batch_release(batch); return DAMACY_OK; } - if (!self || !b) - return DAMACY_INVAL; - if (b != &self->handle) { - log_warn("damacy_release_event: foreign handle (not the active batch)"); - return DAMACY_INVAL; - } - uint16_t s = b->slot_idx; - if (s >= DAMACY_N_BATCH_SLOTS) { - log_warn("damacy_release_event: slot_idx=%u out of range", (unsigned)s); - return DAMACY_INVAL; - } - struct render_job* job = - render_job_pool_for_batch_slot(&self->render_jobs, s); - if (!job) - return DAMACY_INVAL; - - // Push the retained-primary context so cuStreamWaitEvent / cuEventRecord - // land on the right device when the caller is on another thread. - struct ctx_guard cg = { 0 }; - enum damacy_status r = ctx_guard_enter(self, &cg); - if (r != DAMACY_OK) - return r; - scheduler_lock(self->sched); - struct damacy_batch_slot* slot = &self->batch_pool.slots[s]; - if (slot->state != BATCH_HELD) { - log_warn( - "damacy_release_event: slot %u not HELD (state=%d); double release?", - (unsigned)s, - (int)slot->state); - r = DAMACY_INVAL; - goto Done; - } - - // Reuse waits on the caller's event. - if (cuStreamWaitEvent(self->wave_pool.stream_post, (CUevent)event, 0) != - CUDA_SUCCESS) { - // Without the wait, release immediately and report the CUDA error. - release_slot_now(slot, job); - r = DAMACY_CUDA; - goto Done; - } - - release_slot_after_event(slot, job); - r = DAMACY_OK; - -Done: + int stopping = self->stopping; + enum damacy_status status = DAMACY_INVAL; + if (!stopping) + status = self->executor->ops->wait_event(self->executor, event); scheduler_unlock(self->sched); - ctx_guard_exit(&cg); - return r; -} - -// --- batch info / stats --------------------------------------------------- - -void -damacy_batch_info(const struct damacy_batch* b, struct damacy_batch_info* out) -{ - if (!out) - return; - memset(out, 0, sizeof(*out)); - if (!b || !b->d || b->slot_idx >= DAMACY_N_BATCH_SLOTS) - return; - const struct damacy* self = b->d; - const struct damacy_batch_slot* slot = &self->batch_pool.slots[b->slot_idx]; - if (slot->state != BATCH_HELD) - return; - out->device_ptr = slot->dev_ptr; - out->rank = self->batch_pool.rank; - out->dtype = self->cfg.dtype; - out->ready_stream = (void*)self->wave_pool.stream_post; - out->batch_id = slot->batch_id; - for (uint8_t d = 0; d < self->batch_pool.rank; ++d) - out->shape[d] = self->batch_pool.shape[d]; - // shape[0] reflects the actual sample count in the batch. - out->shape[0] = (int64_t)slot->n_samples; + if (stopping && batch->buffer->wait_event) + status = batch->buffer->wait_event(batch->buffer, event); + damacy_batch_release(batch); + return status; } void @@ -200,69 +87,29 @@ damacy_stats_get(const struct damacy* self, struct damacy_stats* out) memset(out, 0, sizeof(*out)); return; } - // scheduler_lock guards every metric_record write; without it the - // struct copy below races every plan/pop_wait update. The - // mutex doesn't change observable state, so the const cast is safe. - struct damacy* m = (struct damacy*)self; - scheduler_lock(m->sched); - *out = m->stats; - out->gpu_bytes_committed = gpu_budget_committed(m->budget); - scheduler_unlock(m->sched); - if (m->array_meta_cache) { - struct prefetch_cache_stats cs; - prefetch_cache_stats_get(m->array_meta_cache, &cs); - out->array_meta.hits = cs.counters.hits; - out->array_meta.misses = cs.counters.misses; - } - if (m->shard_index_cache) { - struct prefetch_cache_stats cs; - prefetch_cache_stats_get(m->shard_index_cache, &cs); - out->shard_index.hits = cs.counters.hits; - out->shard_index.misses = cs.counters.misses; + if (self->stopped) { + *out = self->stats; + return; } - if (m->chunk_layout_cache) { - struct prefetch_cache_stats cs; - prefetch_cache_stats_get(m->chunk_layout_cache, &cs); - out->chunk_layout.hits = cs.counters.hits; - out->chunk_layout.misses = cs.counters.misses; - } - if (m->store_meta_async) { - struct metadata_store_async_latency_stats ls; - metadata_store_async_latency_stats_get(m->store_meta_async, &ls); - out->metadata_latency.ops = ls.ops; - out->metadata_latency.stat_ops = ls.stat_ops; - out->metadata_latency.submit_ops = ls.submit_ops; - out->metadata_latency.active = ls.active; - out->metadata_latency.max_active = ls.max_active; - out->metadata_latency.total_sleep_ns = ls.total_sleep_ns; - out->metadata_latency.max_sleep_ns = ls.max_sleep_ns; - struct metadata_store_async_backend_stats fs; - metadata_store_async_backend_stats_get(m->store_meta_async, &fs); - out->metadata_backend.read_jobs = fs.read_jobs; - out->metadata_backend.read_active = fs.read_active; - out->metadata_backend.read_max_active = fs.read_max_active; - struct metadata_store_async_op_latency_stats os; - metadata_store_async_op_latency_stats_get(m->store_meta_async, &os); - for (unsigned k = 0; k < DAMACY_METADATA_OP_LATENCY_NKINDS; ++k) { - out->metadata_op_latency[k].count = os.kinds[k].count; - out->metadata_op_latency[k].sum_ns = os.kinds[k].sum_ns; - out->metadata_op_latency[k].max_ns = os.kinds[k].max_ns; - for (unsigned b = 0; b < DAMACY_METADATA_OP_LATENCY_NBUCKETS; ++b) - out->metadata_op_latency[k].buckets[b] = os.kinds[k].buckets[b]; - } + struct damacy* mutable = (struct damacy*)self; + scheduler_lock(mutable->sched); + *out = self->stats; + if (!self->stopping) { + mutable->planner->ops->stats(mutable->planner, out); + mutable->executor->ops->stats(mutable->executor, out); } + scheduler_unlock(mutable->sched); } void damacy_stats_reset(struct damacy* self) { - if (!self) + if (!self || self->stopped) return; - // Lock so this reset doesn't race the worker's stat writes (see stats_get). scheduler_lock(self->sched); - stats_init(&self->stats); + if (!self->stopping) { + stats_init(&self->stats); + self->planner->ops->reset_stats(self->planner); + } scheduler_unlock(self->sched); - metadata_store_async_latency_stats_reset(self->store_meta_async); - metadata_store_async_backend_stats_reset(self->store_meta_async); - metadata_store_async_op_latency_stats_reset(self->store_meta_async); } diff --git a/src/damacy_push.c b/src/damacy_push.c index bb872251..9f79e8f8 100644 --- a/src/damacy_push.c +++ b/src/damacy_push.c @@ -1,77 +1,23 @@ -#include "damacy.h" - #include "damacy_internal.h" -static int -sample_aabb_extents_match_cfg(const struct damacy_config* cfg, - const struct damacy_aabb* aabb) -{ - for (uint8_t d = 0; d < cfg->sample_rank; ++d) { - int64_t extent = aabb->dims[d].end - aabb->dims[d].beg; - if (extent != cfg->sample_shape[d]) - return 0; - } - return 1; -} - -// Cfg-only validations; URI / dtype / per-array rank checks surface -// asynchronously through prefetch / plan / pop. -static enum damacy_status -push_one(struct damacy* self, const struct damacy_sample* sample) -{ - if (!sample->uri) - return DAMACY_INVAL; - if (sample->aabb.rank == 0 || sample->aabb.rank > DAMACY_MAX_RANK) - return DAMACY_RANK; - if (sample->aabb.rank != self->cfg.sample_rank) - return DAMACY_RANK; - if (!sample_aabb_extents_match_cfg(&self->cfg, &sample->aabb)) - return DAMACY_INVAL; - - if (lookahead_push_with_sample_seq( - &self->lookahead, sample, self->pushed_samples)) - return DAMACY_OOM; - return DAMACY_OK; -} - struct damacy_push_result damacy_push(struct damacy* self, struct damacy_sample_slice samples) { - struct damacy_push_result r = { .unconsumed = samples, .status = DAMACY_OK }; - if (!self) { - r.status = DAMACY_INVAL; - return r; + struct damacy_push_result result = { .unconsumed = samples, + .status = DAMACY_INVAL }; + if (!self || ((!samples.beg || !samples.end) && samples.beg != samples.end) || + samples.beg > samples.end) + return result; + if (self->stopped) { + result.status = DAMACY_SHUTDOWN; + return result; } - if (samples.beg > samples.end) { - r.status = DAMACY_INVAL; - return r; - } - // No ctx_guard: push touches no CUDA. The scheduler lock serializes - // pushers and pairs pushed_samples with planner consumption of the - // prefetcher ready/error prefix. scheduler_lock(self->sched); - if (self->failed_status != DAMACY_OK) { - r.status = DAMACY_SHUTDOWN; - goto Done; - } - for (const struct damacy_sample* s = samples.beg; s != samples.end; ++s) { - uint64_t unconsumed = - prefetcher_unconsumed_count(self->prefetcher, self->pushed_samples); - if (unconsumed >= self->cfg.lookahead_samples) { - r.unconsumed.beg = s; - r.status = DAMACY_AGAIN; - goto Done; - } - enum damacy_status ps = push_one(self, s); - if (ps != DAMACY_OK) { - r.unconsumed.beg = s; - r.status = ps; - goto Done; - } - self->pushed_samples++; - } - r.unconsumed.beg = samples.end; -Done: + if (self->stopping || self->failed_status != DAMACY_OK) + result.status = DAMACY_SHUTDOWN; + else + result = self->planner->ops->push(self->planner, samples); + scheduler_broadcast(self->sched); scheduler_unlock(self->sched); - return r; + return result; } diff --git a/src/damacy_scheduler.c b/src/damacy_scheduler.c index 486a8fd6..c3fde24b 100644 --- a/src/damacy_scheduler.c +++ b/src/damacy_scheduler.c @@ -1,81 +1,51 @@ -#include "damacy.h" - #include "damacy_internal.h" -#include "nvtx/nvtx.h" -#include "wave/wave_input.h" - -#include -// Drains sealed render jobs into free input_slots, planning a fresh batch -// when no render job has work ready. -static enum damacy_status -kick_input_into_free_slots(struct damacy* self, int* changed) +void +damacy_scheduler_enter(void* arg) { - for (;;) { - int target_job = find_render_job_with_work(&self->render_jobs); - if (target_job < 0) { - int plan_changed = 0; - enum damacy_status s = plan_ready_prefetch(self, &plan_changed); - if (s != DAMACY_OK) - return s; - if (changed && plan_changed) - *changed = 1; - if (plan_changed) - continue; - break; - } + struct damacy* self = arg; + if (self->executor->ops->enter_thread) + self->failed_status = self->executor->ops->enter_thread(self->executor); +} - struct wave_input_reservation t = { 0 }; - enum damacy_status s = - wave_input_reserve(&self->wave_pool, (uint16_t)target_job, &t); - if (s != DAMACY_OK) - return s; - if (!wave_input_reservation_has_slot(&t)) - break; - damacy_nvtx_range_pushf("input/slot%d", - wave_input_reservation_slot_index(&t)); - scheduler_unlock(self->sched); - struct store_submit_result submit = wave_input_submit(&self->wave_pool, &t); - scheduler_lock(self->sched); - s = wave_input_commit(&self->wave_pool, &t, submit, changed); - damacy_nvtx_range_pop(); - if (s != DAMACY_OK) - return s; - if (!any_slot_free(&self->wave_pool)) - break; - } - return DAMACY_OK; +void +damacy_scheduler_leave(void* arg) +{ + struct damacy* self = arg; + if (self->executor->ops->leave_thread) + self->executor->ops->leave_thread(self->executor); } -// One scheduler tick, under scheduler_lock. Lazy ctx push on first call. -// *changed contract (authoritative): every transition site -// (wave_pool_advance, plan_ready_prefetch/plan_commit, -// wave_input_commit) OR-sets it on a real state transition; the worker -// broadcasts iff non-zero. int damacy_scheduler_step(void* arg) { - struct damacy* self = (struct damacy*)arg; - if (!self->worker_ctx_pushed) { - if (self->worker_ctx) - cuCtxPushCurrent(self->worker_ctx); - self->worker_ctx_pushed = 1; - } - self->stats.worker_steps++; - // Wake any pop waiter so it can observe the latched error. - if (self->failed_status != DAMACY_OK) - return 1; - + struct damacy* self = arg; + if (self->stopping || self->failed_status != DAMACY_OK) + return 0; + ++self->stats.worker_steps; int changed = 0; - enum damacy_status r = wave_pool_advance(&self->wave_pool, &changed); - if (r == DAMACY_OK && self->failed_status == DAMACY_OK) - r = kick_input_into_free_slots(self, &changed); - // Retriable backpressure: the reservation was already rolled back, so - // retry next tick rather than latching a fatal error. - if (r == DAMACY_AGAIN) - return changed; - if (r != DAMACY_OK && self->failed_status == DAMACY_OK) { - self->failed_status = r; + enum damacy_status status = + self->executor->ops->step(self->executor, &changed); + if (status == DAMACY_AGAIN) + status = DAMACY_OK; + if (status == DAMACY_OK) + status = pipeline_prepare(self, &changed); + while (status == DAMACY_OK && self->plan_count) { + struct prepared_plan* plan = self->plans[self->plan_head]; + status = + self->executor->ops->submit(self->executor, plan, self->next_batch_id); + if (status == DAMACY_AGAIN) + return changed; + if (status != DAMACY_OK) + break; + self->plans[self->plan_head] = NULL; + self->plan_head = (self->plan_head + 1) % self->queues.prepared_batches; + --self->plan_count; + ++self->next_batch_id; + changed = 1; + } + if (status != DAMACY_OK) { + self->failed_status = status; return 1; } return changed; diff --git a/src/decoder/blosc1_parse.cu b/src/decoder/blosc1_parse.cu index 3eebabe5..88ce97fc 100644 --- a/src/decoder/blosc1_parse.cu +++ b/src/decoder/blosc1_parse.cu @@ -3,7 +3,7 @@ #include "assemble/assemble.h" #include "damacy_limits.h" #include "decoder/launch_check.h" -#include "planner/planner.h" +#include "executor/dispatch.h" #include "zarr/zarr_chunk_layout.h" #include "zarr/zarr_metadata.h" @@ -61,24 +61,31 @@ blosc1_chunk_scan_kernel(const uint8_t* __restrict__ d_compressed, const struct sample_plan& sp = d_sample_plans[chunk.sample_idx_in_batch]; const uint8_t* d_comp = d_compressed + chunk.compressed_offset; + if (chunk.compressed_nbytes < kBloscHeaderBytes) { + record_parse_err(d_parse_err, 4u); + return; + } const uint8_t flags = d_comp[2]; const uint32_t cbytes = read_u32_le(d_comp + 12); - if (cbytes != chunk.compressed_nbytes) { + if (cbytes != chunk.compressed_nbytes || + read_u32_le(d_comp + 4) != chunk.decompressed_nbytes) { record_parse_err(d_parse_err, 4u); return; } - const uint32_t nblocks = sp.layout.nblocks; const bool memcpyed = ((flags >> 1) & 0x1u) != 0u; struct assemble_chunk* a = &d_assemble_chunks[chunk_idx]; if (memcpyed) { + if (cbytes - kBloscHeaderBytes != chunk.decompressed_nbytes) { + record_parse_err(d_parse_err, 4u); + return; + } // Set the per-chunk bit; Kernel B reads it to skip these chunks. atomicOr(&d_is_memcpyed[chunk_idx >> 5], 1u << (chunk_idx & 31u)); - const uint32_t overhead = kBloscHeaderBytes + 4u * nblocks; const uint32_t slot = atomicAdd(d_n_memcpy, 1u); struct gpu_memcpy_op op; - op.d_src = (uint8_t*)d_comp + overhead; + op.d_src = (uint8_t*)d_comp + kBloscHeaderBytes; op.d_dst = d_decompressed + chunk.decompressed_offset; op.nbytes = chunk.decompressed_nbytes; d_memcpy_ops[slot] = op; diff --git a/src/planner/coalesce.c b/src/executor/coalesce.c similarity index 97% rename from src/planner/coalesce.c rename to src/executor/coalesce.c index ca24538c..1cb70981 100644 --- a/src/planner/coalesce.c +++ b/src/executor/coalesce.c @@ -1,9 +1,9 @@ -#include "planner/coalesce.h" +#include "executor/coalesce.h" -#include "planner/read_op_sort.h" +#include "executor/read_op_sort.h" enum damacy_status -coalesce_chunks(struct planner_output* out, +coalesce_chunks(struct dispatch_output* out, uint64_t read_op_max_bytes, uint32_t max_chunks_per_wave, uint32_t* u32_scratch, diff --git a/src/planner/coalesce.h b/src/executor/coalesce.h similarity index 82% rename from src/planner/coalesce.h rename to src/executor/coalesce.h index 137108e0..118e3c66 100644 --- a/src/planner/coalesce.h +++ b/src/executor/coalesce.h @@ -1,7 +1,7 @@ // Coalesce step of the IO planning pipeline: // filter (emit) → sort → fuse-with-cap → interleave → group-by-read. // -// Operates in place on planner_output: sorts the per-chunk +// Operates in place on dispatch_output: sorts the per-chunk // page-aligned read windows by (shard_path, file_offset), then // greedily fuses adjacent windows in the same shard into one // read_op, bounded by read_op_max_bytes. Fused ops are emitted @@ -12,7 +12,7 @@ // read_op. Fill chunk_plans (path empty, nbytes == 0) keep their // 1:1 placeholder read_ops untouched. // -// Populates planner_output.n_chunks_to_load and n_loads_issued as +// Populates dispatch_output.n_chunks_to_load and n_loads_issued as // part of the same pass. // // Cap policy: a single per-chunk read window that already exceeds @@ -21,8 +21,8 @@ // shrinks a single chunk's read either). #pragma once -#include "damacy.h" // damacy_status -#include "planner/planner.h" // planner_output, read_op +#include "damacy.h" // damacy_status +#include "executor/dispatch.h" // dispatch_output, read_op #include @@ -34,7 +34,7 @@ extern "C" // Scratch requirements: // u32_scratch: >= 4 * out->n_read_ops uint32_t slots // read_op_scratch: >= out->n_read_ops slots - enum damacy_status coalesce_chunks(struct planner_output* out, + enum damacy_status coalesce_chunks(struct dispatch_output* out, uint64_t read_op_max_bytes, uint32_t max_chunks_per_wave, uint32_t* u32_scratch, diff --git a/src/executor/cpu_executor.c b/src/executor/cpu_executor.c new file mode 100644 index 00000000..dc5b735f --- /dev/null +++ b/src/executor/cpu_executor.c @@ -0,0 +1,646 @@ +#define ZSTD_STATIC_LINKING_ONLY +#include +#include + +#include "pipeline/components.h" + +#include "damacy_config.h" +#include "damacy_stats.h" +#include "threadpool/threadpool.h" + +#include +#include + +enum cpu_slot_state +{ + CPU_FREE, + CPU_RENDERING, + CPU_READY, + CPU_HELD +}; + +struct cpu_slot +{ + enum cpu_slot_state state; + struct damacy_buffer* buffer; + struct prepared_plan* plan; + uint64_t batch_id; + uint32_t dispatched; + uint32_t remaining; +}; + +struct cpu_worker +{ + void* decoded; + void* workspace; + ZSTD_DCtx* zstd; +}; + +struct cpu_wave +{ + struct store_event event; + struct store_read* reads; + void* input; + enum damacy_status* results; + float* decode_ms; + float* assemble_ms; + uint64_t* output_bytes; + struct platform_clock clock; + uint32_t first_chunk; + uint32_t count; + uint32_t slot; + uint64_t input_bytes; + int active; +}; + +struct cpu_executor +{ + struct damacy_executor base; + struct damacy_reader* reader; + struct damacy_cpu_config config; + struct damacy_batch_spec output; + int64_t shape[DAMACY_MAX_RANK + 1]; + int64_t strides[DAMACY_MAX_RANK + 1]; + struct damacy_stats* stats; + struct threadpool* pool; + struct cpu_worker* workers; + struct cpu_slot slots[2]; + struct cpu_wave waves[2]; + struct cpu_wave* decoding; + uint64_t committed; +}; + +static float +half_to_float(uint16_t half) +{ + uint32_t sign = (uint32_t)(half & 0x8000) << 16; + uint32_t exponent = (half >> 10) & 31; + uint32_t mantissa = half & 1023; + uint32_t bits; + if (!exponent) { + if (!mantissa) + bits = sign; + else { + int shift = 0; + while (!(mantissa & 1024)) { + mantissa <<= 1; + ++shift; + } + bits = sign | (uint32_t)(113 - shift) << 23 | (mantissa & 1023) << 13; + } + } else { + bits = + sign | (exponent == 31 ? 255 : exponent + 112) << 23 | mantissa << 13; + } + float value; + memcpy(&value, &bits, sizeof(value)); + return value; +} + +static uint16_t +float_to_bfloat(float value) +{ + uint32_t bits; + memcpy(&bits, &value, sizeof(bits)); + if ((bits & 0x7fffffff) > 0x7f800000) + return 0x7fff; + return (uint16_t)((bits + 0x7fff + ((bits >> 16) & 1)) >> 16); +} + +static float +source_value(const void* source, enum dtype dtype, uint64_t index) +{ + const unsigned char* bytes = source; + switch (dtype) { + case dtype_u8: + return bytes[index]; + case dtype_u16: { + uint16_t value; + memcpy(&value, bytes + index * sizeof(value), sizeof(value)); + return value; + } + case dtype_i16: { + int16_t value; + memcpy(&value, bytes + index * sizeof(value), sizeof(value)); + return value; + } + case dtype_u32: { + uint32_t value; + memcpy(&value, bytes + index * sizeof(value), sizeof(value)); + return (float)value; + } + case dtype_i32: { + int32_t value; + memcpy(&value, bytes + index * sizeof(value), sizeof(value)); + return (float)value; + } + case dtype_f16: { + uint16_t value; + memcpy(&value, bytes + index * sizeof(value), sizeof(value)); + return half_to_float(value); + } + case dtype_f32: { + float value; + memcpy(&value, bytes + index * sizeof(value), sizeof(value)); + return value; + } + default: + return 0; + } +} + +static uint64_t +assemble_chunk(struct cpu_executor* self, + struct cpu_slot* slot, + const struct plan_chunk* chunk, + const void* decoded) +{ + const struct prepared_plan* plan = slot->plan; + const struct zarr_metadata* meta = &plan->arrays[chunk->array].metadata; + float fill = source_value(meta->fill_value, meta->dtype, 0); + uint64_t output_bytes = 0; + for (uint32_t u = chunk->first_use; u != UINT32_MAX; u = plan->uses[u].next) { + const struct plan_region* region = &plan->regions[plan->uses[u].region]; + uint64_t lo[DAMACY_MAX_RANK], hi[DAMACY_MAX_RANK]; + uint64_t origin[DAMACY_MAX_RANK], coordinate[DAMACY_MAX_RANK]; + uint64_t region_bytes = damacy_dtype_bpe(self->output.dtype); + for (uint8_t d = 0; d < meta->rank; ++d) { + origin[d] = chunk->coordinate[d] * meta->inner_chunk_shape[d]; + uint64_t begin = (uint64_t)region->source.dims[d].beg; + uint64_t end = (uint64_t)region->source.dims[d].end; + lo[d] = origin[d] > begin ? origin[d] : begin; + uint64_t chunk_end = origin[d] + meta->inner_chunk_shape[d]; + hi[d] = chunk_end < end ? chunk_end : end; + coordinate[d] = lo[d]; + region_bytes *= hi[d] - lo[d]; + } + output_bytes += region_bytes; + uint8_t last = meta->rank - 1; + uint64_t width = hi[last] - lo[last]; + for (;;) { + uint64_t source = 0; + uint64_t destination = (uint64_t)region->sample * self->strides[0]; + for (uint8_t d = 0; d < meta->rank; ++d) { + source = + source * meta->inner_chunk_shape[d] + coordinate[d] - origin[d]; + destination += (coordinate[d] - (uint64_t)region->source.dims[d].beg) * + (uint64_t)self->strides[d + 1]; + } + if (!chunk->missing && meta->dtype == dtype_f32 && + self->output.dtype == DAMACY_F32) + memcpy((float*)slot->buffer->data + destination, + (const char*)decoded + source * sizeof(float), + width * sizeof(float)); + else if (self->output.dtype == DAMACY_F32) { + float* dst = (float*)slot->buffer->data + destination; + for (uint64_t j = 0; j < width; ++j) + dst[j] = chunk->missing + ? fill + : source_value(decoded, meta->dtype, source + j); + } else { + uint16_t* dst = (uint16_t*)slot->buffer->data + destination; + for (uint64_t j = 0; j < width; ++j) + dst[j] = float_to_bfloat( + chunk->missing ? fill + : source_value(decoded, meta->dtype, source + j)); + } + int finished = 1; + for (int d = (int)last - 1; d >= 0; --d) { + if (++coordinate[d] < hi[d]) { + finished = 0; + break; + } + coordinate[d] = lo[d]; + } + if (finished) + break; + } + } + return output_bytes; +} + +static enum damacy_status +decode_chunk(struct cpu_worker* worker, + const struct plan_chunk* chunk, + const struct zarr_metadata* metadata, + const void* input, + const void** decoded) +{ + *decoded = worker->decoded; + if (chunk->missing) + return DAMACY_OK; + switch (metadata->inner_codec.id) { + case CODEC_NONE: + if (chunk->encoded_bytes != chunk->decoded_bytes) + return DAMACY_DECODE; + *decoded = input; + return DAMACY_OK; + case CODEC_ZSTD: { + size_t size = ZSTD_decompressDCtx(worker->zstd, + worker->decoded, + chunk->decoded_bytes, + input, + chunk->encoded_bytes); + return !ZSTD_isError(size) && size == chunk->decoded_bytes + ? DAMACY_OK + : DAMACY_DECODE; + } + case CODEC_BLOSC_ZSTD: { + size_t size = 0; + if (blosc_cbuffer_validate(input, chunk->encoded_bytes, &size) || + size != chunk->decoded_bytes || + chunk->encoded_bytes < BLOSC_MIN_HEADER_LENGTH) + return DAMACY_DECODE; + size_t nbytes, cbytes, blocksize; + blosc_cbuffer_sizes(input, &nbytes, &cbytes, &blocksize); + size_t typesize; + int flags; + blosc_cbuffer_metainfo(input, &typesize, &flags); + const char* compressor = blosc_cbuffer_complib(input); + if (!blocksize || blocksize > size || + typesize != dtype_bpe(metadata->dtype) || + cbytes != chunk->encoded_bytes || size > BLOSC_MAX_BUFFERSIZE || + (!(flags & BLOSC_MEMCPYED) && + (!compressor || strcmp(compressor, BLOSC_ZSTD_LIBNAME)))) + return DAMACY_DECODE; + int result = blosc_decompress_ctx(input, worker->decoded, size, 1); + return result > 0 && (size_t)result == size ? DAMACY_OK : DAMACY_DECODE; + } + default: + return DAMACY_DECODE; + } +} + +static void +decode_one(size_t index, int tid, void* arg) +{ + struct cpu_executor* self = arg; + struct cpu_wave* wave = self->decoding; + struct cpu_slot* slot = &self->slots[wave->slot]; + const struct plan_chunk* chunk = + &slot->plan->chunks[wave->first_chunk + index]; + const struct zarr_metadata* meta = &slot->plan->arrays[chunk->array].metadata; + const void* input = + (const char*)wave->input + index * self->config.max_encoded_chunk_bytes; + const void* decoded; + struct platform_clock clock = { 0 }; + platform_toc(&clock); + wave->results[index] = + decode_chunk(&self->workers[tid], chunk, meta, input, &decoded); + wave->decode_ms[index] = platform_toc(&clock) * 1000; + wave->assemble_ms[index] = 0; + wave->output_bytes[index] = 0; + if (wave->results[index] == DAMACY_OK) { + wave->output_bytes[index] = assemble_chunk(self, slot, chunk, decoded); + wave->assemble_ms[index] = platform_toc(&clock) * 1000; + } +} + +static void +cpu_buffer_destroy(struct damacy_buffer* buffer) +{ + free(buffer->data); + free(buffer); +} + +static void +cpu_stop(struct damacy_executor* base) +{ + struct cpu_executor* self = (void*)base; + for (unsigned i = 0; i < 2; ++i) { + struct cpu_wave* wave = &self->waves[i]; + if (wave->active && wave->event.impl) + store_event_wait(self->reader->store, wave->event); + free(wave->input); + free(wave->reads); + free(wave->results); + free(wave->decode_ms); + free(wave->assemble_ms); + free(wave->output_bytes); + *wave = (struct cpu_wave){ 0 }; + prepared_plan_destroy(self->slots[i].plan); + buffer_release(self->slots[i].buffer); + self->slots[i] = (struct cpu_slot){ 0 }; + } + threadpool_free(self->pool); + self->pool = NULL; + if (self->workers) { + for (uint32_t i = 0; i < self->config.decode_workers; ++i) { + free(self->workers[i].decoded); + free(self->workers[i].workspace); + } + free(self->workers); + self->workers = NULL; + } + self->committed = 0; +} + +static enum damacy_status +cpu_start(struct damacy_executor* base, + const struct damacy_batch_spec* output, + struct damacy_stats* stats) +{ + struct cpu_executor* self = (void*)base; + self->output = *output; + self->stats = stats; + uint64_t bytes; + enum damacy_status status = + batch_spec_layout(output, self->shape, self->strides, &bytes); + if (status != DAMACY_OK) + return status; + uint64_t workers = self->config.decode_workers; + uint64_t workspace = ZSTD_estimateDCtxSize(); + uint64_t codec_reserve = + workspace + 3ull * self->config.max_decoded_chunk_bytes + (256u << 10); + uint64_t per_worker = + 2ull * self->config.max_encoded_chunk_bytes + + self->config.max_decoded_chunk_bytes + codec_reserve + + sizeof(struct cpu_worker) + + 2 * (sizeof(struct store_read) + sizeof(enum damacy_status) + + 2 * sizeof(float) + sizeof(uint64_t)); + uint64_t fixed = sizeof(*self) + 2 * sizeof(struct damacy_buffer); + if (per_worker > self->config.max_memory_bytes / workers || + bytes > self->config.max_memory_bytes / 2) + return DAMACY_BUDGET; + uint64_t need = fixed + per_worker * workers + 2 * bytes; + if (need < bytes || need > self->config.max_memory_bytes) + return DAMACY_BUDGET; + self->committed = need; + for (unsigned i = 0; i < 2; ++i) { + struct damacy_buffer* buffer = calloc(1, sizeof(*buffer)); + if (!buffer) + goto Fail; + atomic_init(&buffer->references, 1); + buffer->destroy = cpu_buffer_destroy; + buffer->device_type = DAMACY_DEVICE_CPU; + buffer->nbytes = bytes; + self->slots[i].buffer = buffer; + buffer->data = malloc((size_t)bytes); + if (!buffer->data) + goto Fail; + struct cpu_wave* wave = &self->waves[i]; + wave->input = + malloc((size_t)workers * self->config.max_encoded_chunk_bytes); + wave->reads = calloc((size_t)workers, sizeof(*wave->reads)); + wave->results = calloc((size_t)workers, sizeof(*wave->results)); + wave->decode_ms = calloc((size_t)workers, sizeof(*wave->decode_ms)); + wave->assemble_ms = calloc((size_t)workers, sizeof(*wave->assemble_ms)); + wave->output_bytes = calloc((size_t)workers, sizeof(*wave->output_bytes)); + if (!wave->input || !wave->reads || !wave->results || !wave->decode_ms || + !wave->assemble_ms || !wave->output_bytes) + goto Fail; + } + self->workers = calloc((size_t)workers, sizeof(*self->workers)); + if (!self->workers) + goto Fail; + for (uint32_t i = 0; i < workers; ++i) { + struct cpu_worker* worker = &self->workers[i]; + worker->decoded = malloc(self->config.max_decoded_chunk_bytes); + worker->workspace = malloc((size_t)workspace); + if (!worker->decoded || !worker->workspace) + goto Fail; + worker->zstd = ZSTD_initStaticDCtx(worker->workspace, (size_t)workspace); + if (!worker->zstd) + goto Fail; + } + self->pool = threadpool_new((int)workers - 1); + if (!self->pool) + goto Fail; + return DAMACY_OK; +Fail: + cpu_stop(base); + return DAMACY_OOM; +} + +static enum damacy_status +cpu_submit(struct damacy_executor* base, + struct prepared_plan* plan, + uint64_t batch_id) +{ + struct cpu_executor* self = (void*)base; + int slot = -1; + for (unsigned i = 0; i < 2; ++i) + if (self->slots[i].state == CPU_FREE) { + slot = (int)i; + break; + } + if (slot < 0) + return DAMACY_AGAIN; + for (uint32_t i = 0; i < plan->n_chunks; ++i) { + const struct plan_chunk* chunk = &plan->chunks[i]; + if (chunk->encoded_bytes > self->config.max_encoded_chunk_bytes || + chunk->decoded_bytes > self->config.max_decoded_chunk_bytes) + return DAMACY_BUDGET; + enum compression_codec codec = + plan->arrays[chunk->array].metadata.inner_codec.id; + if (!chunk->missing && codec != CODEC_NONE && codec != CODEC_ZSTD && + codec != CODEC_BLOSC_ZSTD) + return DAMACY_DECODE; + } + struct cpu_slot* target = &self->slots[slot]; + target->plan = plan; + target->batch_id = batch_id; + target->dispatched = 0; + target->remaining = plan->n_chunks; + target->state = CPU_RENDERING; + return DAMACY_OK; +} + +static enum damacy_status +cpu_dispatch(struct cpu_executor* self, struct cpu_wave* wave, int* changed) +{ + int slot_index = -1; + for (unsigned i = 0; i < 2; ++i) { + const struct cpu_slot* slot = &self->slots[i]; + if (slot->state == CPU_RENDERING && + slot->dispatched < slot->plan->n_chunks && + (slot_index < 0 || slot->batch_id < self->slots[slot_index].batch_id)) + slot_index = (int)i; + } + if (slot_index < 0) + return DAMACY_OK; + struct cpu_slot* slot = &self->slots[slot_index]; + uint32_t count = slot->plan->n_chunks - slot->dispatched; + if (count > self->config.decode_workers) + count = self->config.decode_workers; + if (count > self->reader->max_inflight_reads) + count = self->reader->max_inflight_reads; + uint32_t n_reads = 0; + uint64_t bytes = 0; + for (uint32_t i = 0; i < count; ++i) { + const struct plan_chunk* chunk = &slot->plan->chunks[slot->dispatched + i]; + if (chunk->missing) + continue; + wave->reads[n_reads++] = + (struct store_read){ .key = chunk->path, + .offset = chunk->offset, + .len = chunk->encoded_bytes, + .dst = + (char*)wave->input + + (size_t)i * self->config.max_encoded_chunk_bytes }; + bytes += chunk->encoded_bytes; + } + platform_toc(&wave->clock); + struct store_submit_result result = + store_read_submit(self->reader->store, wave->reads, n_reads); + if (result.status != DAMACY_OK) + return result.status; + wave->event = result.event; + wave->active = 1; + wave->slot = (uint32_t)slot_index; + wave->first_chunk = slot->dispatched; + wave->count = count; + wave->input_bytes = bytes; + slot->dispatched += count; + self->stats->reads_issued += n_reads; + self->stats->chunks_dispatched += count; + ++self->stats->waves_emitted; + *changed = 1; + return DAMACY_OK; +} + +static enum damacy_status +cpu_step(struct damacy_executor* base, int* changed) +{ + struct cpu_executor* self = (void*)base; + for (unsigned i = 0; i < 2; ++i) { + struct cpu_slot* slot = &self->slots[i]; + if (slot->state == CPU_HELD && buffer_available(slot->buffer)) { + slot->state = CPU_FREE; + *changed = 1; + } + } + for (unsigned i = 0; i < 2; ++i) { + if (!self->waves[i].active) { + enum damacy_status status = cpu_dispatch(self, &self->waves[i], changed); + if (status != DAMACY_OK && status != DAMACY_AGAIN) + return status; + } + } + for (unsigned i = 0; i < 2; ++i) { + struct cpu_wave* wave = &self->waves[i]; + if (!wave->active) + continue; + struct store_event_poll poll = + store_event_query(self->reader->store, wave->event); + if (!poll.ready) + continue; + wave->event = (struct store_event){ 0 }; + if (poll.status != DAMACY_OK) + return poll.status; + metric_record(&self->stats->io, + platform_toc(&wave->clock) * 1000, + (double)wave->input_bytes, + (double)wave->input_bytes); + self->decoding = wave; + threadpool_for_n_dynamic(self->pool, wave->count, decode_one, self); + struct cpu_slot* slot = &self->slots[wave->slot]; + for (uint32_t j = 0; j < wave->count; ++j) { + if (wave->results[j] != DAMACY_OK) + return wave->results[j]; + const struct plan_chunk* chunk = + &slot->plan->chunks[wave->first_chunk + j]; + metric_record(&self->stats->decode, + wave->decode_ms[j], + chunk->encoded_bytes, + chunk->missing ? 0 : chunk->decoded_bytes); + metric_record(&self->stats->assemble, + wave->assemble_ms[j], + chunk->decoded_bytes, + wave->output_bytes[j]); + } + slot->remaining -= wave->count; + wave->active = 0; + if (!slot->remaining) { + slot->state = CPU_READY; + prepared_plan_destroy(slot->plan); + slot->plan = NULL; + } + *changed = 1; + } + return DAMACY_OK; +} + +static enum damacy_status +cpu_take(struct damacy_executor* base, struct damacy_batch** out) +{ + struct cpu_executor* self = (void*)base; + int oldest = -1; + for (unsigned i = 0; i < 2; ++i) { + struct cpu_slot* slot = &self->slots[i]; + if ((slot->state == CPU_RENDERING || slot->state == CPU_READY) && + (oldest < 0 || slot->batch_id < self->slots[oldest].batch_id)) + oldest = (int)i; + } + if (oldest < 0 || self->slots[oldest].state != CPU_READY) + return DAMACY_AGAIN; + struct cpu_slot* slot = &self->slots[oldest]; + *out = batch_create(slot->buffer, &self->output, slot->batch_id); + if (!*out) + return DAMACY_OOM; + slot->state = CPU_HELD; + return DAMACY_OK; +} + +static enum damacy_status +cpu_wait_event(struct damacy_executor* base, void* event) +{ + (void)base; + (void)event; + return DAMACY_INVAL; +} + +static int +cpu_busy(const struct damacy_executor* base) +{ + const struct cpu_executor* self = (const void*)base; + return self->slots[0].state != CPU_FREE || self->slots[1].state != CPU_FREE; +} + +static void +cpu_stats(struct damacy_executor* base, struct damacy_stats* out) +{ + out->host_bytes_committed = ((struct cpu_executor*)base)->committed; +} + +static void +cpu_destroy(struct damacy_executor* base) +{ + cpu_stop(base); + free(base); +} + +static const struct damacy_executor_ops cpu_ops = { .start = cpu_start, + .submit = cpu_submit, + .step = cpu_step, + .take = cpu_take, + .wait_event = + cpu_wait_event, + .busy = cpu_busy, + .stats = cpu_stats, + .stop = cpu_stop, + .destroy = cpu_destroy }; + +enum damacy_status +damacy_cpu_executor_create(struct damacy_reader* reader, + const struct damacy_cpu_config* config, + struct damacy_executor** out) +{ + if (!out) + return DAMACY_INVAL; + *out = NULL; + if (!reader || !config || !config->decode_workers || + config->decode_workers > DAMACY_MAX_IO_THREADS || + !config->max_encoded_chunk_bytes || !config->max_decoded_chunk_bytes || + !config->max_memory_bytes) + return DAMACY_INVAL; + struct cpu_executor* self = calloc(1, sizeof(*self)); + if (!self) + return DAMACY_OOM; + self->base.ops = &cpu_ops; + self->base.device_type = DAMACY_DEVICE_CPU; + self->reader = reader; + self->config = *config; + *out = &self->base; + return DAMACY_OK; +} diff --git a/src/executor/cuda_executor.c b/src/executor/cuda_executor.c new file mode 100644 index 00000000..ac20f70b --- /dev/null +++ b/src/executor/cuda_executor.c @@ -0,0 +1,623 @@ +#include "executor/cuda_executor.h" +#include "pipeline/components.h" + +#include "batch_pool/batch_pool.h" +#include "damacy_config.h" +#include "executor/dispatch.h" +#include "gpu_budget/gpu_budget.h" +#include "log/log.h" +#include "numa/numa_cuda.h" +#include "render_job/render_job.h" +#include "store/store_fs_gds.h" +#include "util/cuda_check.h" +#include "wave/wave_budget.h" +#include "wave/wave_input.h" +#include "wave/wave_pool.h" + +#include +#include +#include + +#include "cuda_geometry.inc" + +struct cuda_context +{ + _Atomic uint32_t references; + CUcontext handle; + CUstream retained_stream; + int retained_device; +}; + +struct cuda_layout_entry +{ + char* uri; + struct chunk_layout layout; +}; + +struct cuda_executor +{ + struct damacy_executor base; + struct damacy_reader* reader; + struct damacy_cuda_config config; + struct damacy_batch_spec output; + struct damacy_config cfg; + struct cuda_context* context; + struct numa_resolved numa; + struct gpu_budget* budget; + struct store* gds; + struct damacy_stats* stats; + struct damacy_batch_pool batches; + struct render_job_pool jobs; + struct wave_pool waves; + struct dispatch_scratch scratch; + struct prepared_plan* plans[DAMACY_N_BATCH_SLOTS]; + struct damacy_buffer* buffers[DAMACY_N_BATCH_SLOTS]; + struct cuda_layout_entry* layouts; + uint32_t next_layout; + uint64_t max_read_bytes; + int worker_context; +}; + +static void +context_release(struct cuda_context* context) +{ + if (context && atomic_fetch_sub_explicit( + &context->references, 1, memory_order_acq_rel) == 1) { + if (context->retained_stream && + cuCtxPushCurrent(context->handle) == CUDA_SUCCESS) { + cuStreamDestroy(context->retained_stream); + cuCtxPopCurrent(NULL); + } + if (context->retained_device >= 0) + cuDevicePrimaryCtxRelease((CUdevice)context->retained_device); + free(context); + } +} + +static void +cuda_buffer_destroy(struct damacy_buffer* buffer) +{ + struct cuda_context* context = buffer->context; + if (cuCtxPushCurrent(context->handle) == CUDA_SUCCESS) { + cuMemFree(CUDPTR(buffer->data)); + cuCtxPopCurrent(NULL); + } + context_release(context); + free(buffer); +} + +static enum damacy_status +cuda_buffer_wait_event(struct damacy_buffer* buffer, void* event) +{ + struct cuda_context* context = buffer->context; + if (cuCtxPushCurrent(context->handle) != CUDA_SUCCESS) + return DAMACY_CUDA; + CUresult result = cuEventSynchronize(event); + cuCtxPopCurrent(NULL); + return result == CUDA_SUCCESS ? DAMACY_OK : DAMACY_CUDA; +} + +static void +cuda_stop(struct damacy_executor* base) +{ + struct cuda_executor* self = (void*)base; + int pushed = + self->context && cuCtxPushCurrent(self->context->handle) == CUDA_SUCCESS; + if (pushed && self->waves.stream_post) { + self->context->retained_stream = self->waves.stream_post; + cuStreamSynchronize(self->waves.stream_post); + self->waves.stream_post = NULL; + } + wave_pool_destroy(&self->waves, !pushed); + render_job_pool_destroy(&self->jobs, !pushed); + for (unsigned i = 0; i < DAMACY_N_BATCH_SLOTS; ++i) { + prepared_plan_destroy(self->plans[i]); + self->plans[i] = NULL; + if (self->buffers[i]) { + self->batches.slots[i].dev_ptr = NULL; + buffer_release(self->buffers[i]); + self->buffers[i] = NULL; + } + } + batch_pool_destroy(&self->batches, !pushed); + gpu_budget_destroy(self->budget); + self->budget = NULL; + store_destroy(self->gds); + self->gds = NULL; + dispatch_scratch_destroy(&self->scratch); + if (self->layouts) { + for (unsigned i = 0; i < self->config.chunk_layout_entries; ++i) + free(self->layouts[i].uri); + free(self->layouts); + self->layouts = NULL; + } + if (pushed) + cuCtxPopCurrent(NULL); + context_release(self->context); + self->context = NULL; +} + +static enum damacy_status +cuda_start(struct damacy_executor* base, + const struct damacy_batch_spec* output, + struct damacy_stats* stats) +{ + struct cuda_executor* self = (void*)base; + self->output = *output; + self->stats = stats; + self->cfg = (struct damacy_config){ + .dtype = output->dtype, + .sample_rank = output->sample_rank, + .samples_per_batch = output->samples_per_batch, + .device = self->config.device, + .tuning = { .max_gpu_memory_bytes = self->config.max_gpu_memory_bytes, + .max_chunk_uncompressed_bytes = self->config.max_chunk_bytes, + .max_read_op_bytes = self->config.max_read_bytes, + .host_buffer_waves = self->config.host_buffer_waves, + .max_chunks_per_wave = self->config.max_chunks_per_wave, + .max_substreams_per_chunk = + self->config.max_substreams_per_chunk, + .numa_strategy = self->config.numa_strategy, + .numa_node = self->config.numa_node, + .enable_gds = self->config.enable_gds }, + .debug = { .bypass_decode = self->config.bypass_decode } + }; + if (self->cfg.tuning.max_chunks_per_wave > self->reader->max_inflight_reads) + self->cfg.tuning.max_chunks_per_wave = self->reader->max_inflight_reads; + memcpy(self->cfg.sample_shape, + output->sample_shape, + sizeof(self->cfg.sample_shape)); + enum damacy_status status = DAMACY_CUDA; + int pushed = 0; + if (cuInit(0) != CUDA_SUCCESS) + return status; + CUcontext caller = NULL; + if (cuCtxGetCurrent(&caller) != CUDA_SUCCESS) + return status; + self->context = calloc(1, sizeof(*self->context)); + if (!self->context) + return DAMACY_OOM; + atomic_init(&self->context->references, 1); + self->context->retained_device = -1; + CUdevice device; + if (self->config.device >= 0) { + if (caller) { + if (cuCtxGetDevice(&device) != CUDA_SUCCESS) + goto Fail; + if (device != self->config.device) { + status = DAMACY_INVAL; + goto Fail; + } + } + if (cuDeviceGet(&device, self->config.device) != CUDA_SUCCESS || + cuDevicePrimaryCtxRetain(&self->context->handle, device) != + CUDA_SUCCESS) + goto Fail; + self->context->retained_device = self->config.device; + } else { + if (!caller) { + status = DAMACY_INVAL; + goto Fail; + } + if (cuCtxGetDevice(&device) != CUDA_SUCCESS) + goto Fail; + self->context->handle = caller; + } + self->base.device_id = device; + if (cuCtxPushCurrent(self->context->handle) != CUDA_SUCCESS) + goto Fail; + pushed = 1; + numa_init( + self->config.numa_strategy, self->config.numa_node, device, &self->numa); + status = batch_pool_compute_layout(&self->batches, + output->sample_shape, + output->sample_rank, + output->samples_per_batch, + damacy_dtype_bpe(output->dtype)); + if (status != DAMACY_OK) + goto Fail; + if (self->batches.n_bytes >= self->config.max_gpu_memory_bytes / 2) { + status = DAMACY_BUDGET; + goto Fail; + } + uint64_t resolver_budget = + self->config.max_gpu_memory_bytes - 2 * self->batches.n_bytes; + struct resolved_wave_geometry geometry; + status = resolve_wave_geometry( + &self->cfg, resolver_budget, self->config.max_chunk_bytes, &geometry, NULL); + if (status != DAMACY_OK) + goto Fail; + self->max_read_bytes = self->config.max_read_bytes; + if (self->max_read_bytes > geometry.sizing.input_staging_per_wave) + self->max_read_bytes = geometry.sizing.input_staging_per_wave; + status = DAMACY_OOM; + self->budget = gpu_budget_new(self->config.max_gpu_memory_bytes); + if (!self->budget) + goto Fail; + gpu_budget_commit(self->budget, geometry.predicted.total); + if (geometry.want_gds) { + self->gds = + store_fs_gds_create(&(struct store_fs_gds_config){ .root = "" }); + if (!self->gds) { + status = DAMACY_INVAL; + goto Fail; + } + } + struct platform_cpu_mask saved; + numa_scope_enter(&self->numa, &saved); + int failed = 0; + for (unsigned i = 0; i < DAMACY_N_BATCH_SLOTS && !failed; ++i) + failed = render_job_init(&self->jobs.jobs[i], output->samples_per_batch); + if (!failed) + failed = wave_pool_init(&self->waves, + &self->batches, + &self->jobs, + geometry.want_gds ? self->gds : self->reader->store, + stats, + output->dtype, + geometry.host_buffer_waves, + geometry.max_chunks_per_wave, + geometry.max_substreams_per_chunk, + geometry.sizing.input_staging_per_wave, + geometry.sizing.dev_decompressed_per_wave, + self->config.max_chunk_bytes, + geometry.input, + self->config.bypass_decode, + self->budget); + numa_scope_exit(&saved); + if (failed) + goto Fail; + self->layouts = + calloc(self->config.chunk_layout_entries, sizeof(*self->layouts)); + if (!self->layouts) + goto Fail; + cuCtxPopCurrent(NULL); + return DAMACY_OK; +Fail: + if (pushed) + cuCtxPopCurrent(NULL); + cuda_stop(base); + return status; +} + +static enum damacy_status +allocate_outputs(struct cuda_executor* self) +{ + if (self->batches.allocated) + return DAMACY_OK; + uint64_t bytes = 2 * self->batches.n_bytes; + enum damacy_status status = + gpu_budget_try_commit(self->budget, bytes, "batch-output pool"); + if (status != DAMACY_OK) + return status; + status = batch_pool_alloc_dev(&self->batches); + if (status != DAMACY_OK) { + gpu_budget_release(self->budget, bytes); + return status; + } + for (unsigned i = 0; i < DAMACY_N_BATCH_SLOTS; ++i) { + struct damacy_buffer* buffer = calloc(1, sizeof(*buffer)); + if (!buffer) + return DAMACY_OOM; + atomic_init(&buffer->references, 1); + buffer->data = self->batches.slots[i].dev_ptr; + buffer->nbytes = self->batches.n_bytes; + buffer->ready_stream = self->waves.stream_post; + buffer->device_type = DAMACY_DEVICE_CUDA; + buffer->device_id = self->base.device_id; + buffer->destroy = cuda_buffer_destroy; + buffer->wait_event = cuda_buffer_wait_event; + buffer->context = self->context; + atomic_fetch_add_explicit( + &self->context->references, 1, memory_order_relaxed); + self->buffers[i] = buffer; + } + return DAMACY_OK; +} + +static struct cuda_layout_entry* +find_layout(struct cuda_executor* self, const char* uri) +{ + for (unsigned i = 0; i < self->config.chunk_layout_entries; ++i) + if (self->layouts[i].uri && !strcmp(uri, self->layouts[i].uri)) + return &self->layouts[i]; + return NULL; +} + +static enum damacy_status +prepare_layouts(struct cuda_executor* self, + const struct prepared_plan* plan, + struct dispatch_output* output) +{ + for (uint32_t a = 0; a < plan->n_arrays; ++a) { + const struct plan_array* array = &plan->arrays[a]; + enum compression_codec codec = array->metadata.inner_codec.id; + if (codec == CODEC_NONE || codec == CODEC_ZSTD) + continue; + if (codec != CODEC_BLOSC_ZSTD) + return DAMACY_DECODE; + const struct plan_chunk* chunk = NULL; + for (uint32_t i = 0; i < plan->n_chunks; ++i) + if (plan->chunks[i].array == a && !plan->chunks[i].missing) { + chunk = &plan->chunks[i]; + break; + } + if (!chunk) + continue; + struct cuda_layout_entry* entry = find_layout(self, array->uri); + if (!entry) { + ++self->stats->chunk_layout.misses; + if (chunk->encoded_bytes < 16) + return DAMACY_DECODE; + unsigned char header[16]; + struct store_read read = { .key = chunk->path, + .offset = chunk->offset, + .dst = header, + .len = sizeof(header) }; + struct store_submit_result result = + store_read_submit(self->reader->store, &read, 1); + if (result.status != DAMACY_OK) + return result.status; + enum damacy_status status = + store_event_wait(self->reader->store, result.event); + if (status != DAMACY_OK) + return status; + struct chunk_layout layout; + if (zarr_chunk_layout_parse_header(header, + chunk->encoded_bytes, + (uint8_t)codec, + self->config.max_substreams_per_chunk, + &layout) || + layout.nbytes != chunk->decoded_bytes || + layout.typesize != dtype_bpe(array->metadata.dtype)) + return DAMACY_DECODE; + char* uri = strdup(array->uri); + if (!uri) + return DAMACY_OOM; + entry = + &self->layouts[self->next_layout++ % self->config.chunk_layout_entries]; + free(entry->uri); + *entry = (struct cuda_layout_entry){ .uri = uri, .layout = layout }; + } else { + ++self->stats->chunk_layout.hits; + } + for (uint32_t i = 0; i < plan->n_regions; ++i) + if (plan->regions[i].array == a) { + struct sample_plan* sample = + &output->sample_plans[plan->regions[i].sample]; + sample->layout = entry->layout; + sample->layout_probed = 1; + } + } + return DAMACY_OK; +} + +static enum damacy_status +cuda_submit(struct damacy_executor* base, + struct prepared_plan* plan, + uint64_t batch_id) +{ + struct cuda_executor* self = (void*)base; + int slot_index = find_free_batch_slot(&self->batches); + if (slot_index < 0 || find_render_job_with_work(&self->jobs) >= 0) + return DAMACY_AGAIN; + if (cuCtxPushCurrent(self->context->handle) != CUDA_SUCCESS) + return DAMACY_CUDA; + enum damacy_status status = allocate_outputs(self); + if (status != DAMACY_OK) + goto Done; + struct render_job* job = &self->jobs.jobs[slot_index]; + struct dispatch_output output = + render_job_dispatch_output(job, self->output.samples_per_batch); + status = dispatch_plan_build(plan, + (uint16_t)slot_index, + platform_page_alignment(), + self->max_read_bytes, + self->cfg.tuning.max_chunks_per_wave, + &output, + &self->scratch); + if (status != DAMACY_OK) + goto Done; + status = prepare_layouts(self, plan, &output); + if (status != DAMACY_OK) + goto Done; + render_job_commit_plan(job, (uint16_t)slot_index, batch_id, &output); + status = render_job_upload_sample_plans(job, self->waves.stream_input); + if (status != DAMACY_OK) + goto Done; + struct damacy_batch_slot* slot = &self->batches.slots[slot_index]; + slot->batch_id = batch_id; + slot->n_samples = self->output.samples_per_batch; + slot->n_chunks = job->n_chunks; + slot->chunks_remaining = (int32_t)job->n_chunks; + slot->state = BATCH_RENDERING; + self->plans[slot_index] = plan; + self->stats->reads_issued += job->n_loads_issued; +Done: + cuCtxPopCurrent(NULL); + return status; +} + +static enum damacy_status +cuda_enter_thread(struct damacy_executor* base) +{ + struct cuda_executor* self = (void*)base; + if (cuCtxPushCurrent(self->context->handle) != CUDA_SUCCESS) + return DAMACY_CUDA; + self->worker_context = 1; + numa_apply_thread_affinity(&self->numa, "cuda_executor"); + return DAMACY_OK; +} + +static void +cuda_leave_thread(struct damacy_executor* base) +{ + struct cuda_executor* self = (void*)base; + if (self->worker_context) { + cuCtxPopCurrent(NULL); + self->worker_context = 0; + } +} + +static enum damacy_status +cuda_step(struct damacy_executor* base, int* changed) +{ + struct cuda_executor* self = (void*)base; + for (unsigned i = 0; i < DAMACY_N_BATCH_SLOTS; ++i) { + struct damacy_batch_slot* slot = &self->batches.slots[i]; + if (slot->state == BATCH_HELD && buffer_available(self->buffers[i])) { + batch_slot_reset_for_reuse(slot); + render_job_reset(&self->jobs.jobs[i]); + *changed = 1; + } + } + enum damacy_status status = wave_pool_advance(&self->waves, changed); + while (status == DAMACY_OK) { + int index = find_render_job_with_work(&self->jobs); + if (index < 0) + break; + struct wave_input_reservation reservation = { 0 }; + status = wave_input_reserve(&self->waves, (uint16_t)index, &reservation); + if (status != DAMACY_OK || !wave_input_reservation_has_slot(&reservation)) + break; + struct store_submit_result result = + wave_input_submit(&self->waves, &reservation); + status = wave_input_commit(&self->waves, &reservation, result, changed); + if (!any_slot_free(&self->waves)) + break; + } + for (unsigned i = 0; i < DAMACY_N_BATCH_SLOTS; ++i) + if (self->batches.slots[i].state == BATCH_READY && self->plans[i]) { + prepared_plan_destroy(self->plans[i]); + self->plans[i] = NULL; + } + return status; +} + +static enum damacy_status +cuda_take(struct damacy_executor* base, struct damacy_batch** out) +{ + struct cuda_executor* self = (void*)base; + int oldest = -1; + for (unsigned i = 0; i < DAMACY_N_BATCH_SLOTS; ++i) { + const struct damacy_batch_slot* slot = &self->batches.slots[i]; + if ((slot->state == BATCH_RENDERING || slot->state == BATCH_READY) && + (oldest < 0 || slot->batch_id < self->batches.slots[oldest].batch_id)) + oldest = (int)i; + } + if (oldest < 0 || self->batches.slots[oldest].state != BATCH_READY) + return DAMACY_AGAIN; + *out = batch_create( + self->buffers[oldest], &self->output, self->batches.slots[oldest].batch_id); + if (!*out) + return DAMACY_OOM; + self->batches.slots[oldest].state = BATCH_HELD; + return DAMACY_OK; +} + +static enum damacy_status +cuda_wait_event(struct damacy_executor* base, void* event) +{ + struct cuda_executor* self = (void*)base; + if (cuCtxPushCurrent(self->context->handle) != CUDA_SUCCESS) + return DAMACY_CUDA; + CUresult result = cuStreamWaitEvent(self->waves.stream_post, event, 0); + cuCtxPopCurrent(NULL); + return result == CUDA_SUCCESS ? DAMACY_OK : DAMACY_CUDA; +} + +static int +cuda_busy(const struct damacy_executor* base) +{ + const struct cuda_executor* self = (const void*)base; + return any_batch_in_flight(&self->batches); +} + +static void +cuda_stats(struct damacy_executor* base, struct damacy_stats* out) +{ + out->gpu_bytes_committed = + gpu_budget_committed(((struct cuda_executor*)base)->budget); +} + +static void +cuda_destroy(struct damacy_executor* base) +{ + cuda_stop(base); + free(base); +} + +static const struct damacy_executor_ops cuda_ops = { + .enter_thread = cuda_enter_thread, + .leave_thread = cuda_leave_thread, + .start = cuda_start, + .submit = cuda_submit, + .step = cuda_step, + .take = cuda_take, + .wait_event = cuda_wait_event, + .busy = cuda_busy, + .stats = cuda_stats, + .stop = cuda_stop, + .destroy = cuda_destroy +}; + +uint64_t +cuda_executor_set_budget(struct damacy_executor* base, uint64_t value) +{ + return base->ops == &cuda_ops + ? gpu_budget_set_committed_for_test( + ((struct cuda_executor*)base)->budget, value) + : 0; +} + +enum damacy_status +damacy_cuda_executor_create(struct damacy_reader* reader, + const struct damacy_cuda_config* config, + struct damacy_executor** out) +{ + if (!out) + return DAMACY_INVAL; + *out = NULL; + if (!reader || !config || config->device < -1 || + !config->max_gpu_memory_bytes || !config->chunk_layout_entries || + !config->max_chunk_bytes || !config->max_read_bytes || + config->max_read_bytes > UINT32_MAX || !config->max_chunks_per_wave || + config->max_chunks_per_wave > DAMACY_HARD_MAX_CHUNKS_PER_WAVE || + !config->max_substreams_per_chunk || + config->max_substreams_per_chunk > DAMACY_HARD_MAX_SUBSTREAMS_PER_CHUNK || + config->host_buffer_waves < DAMACY_N_WAVES || + config->host_buffer_waves > DAMACY_MAX_HOST_BUFFER_WAVES || + config->numa_strategy < DAMACY_NUMA_AUTO || + config->numa_strategy > DAMACY_NUMA_PIN_TO || + (config->numa_strategy == DAMACY_NUMA_PIN_TO && config->numa_node < 0) || + config->enable_gds < DAMACY_GDS_AUTO || + config->enable_gds > DAMACY_GDS_OFF) + return DAMACY_INVAL; + struct cuda_executor* self = calloc(1, sizeof(*self)); + if (!self) + return DAMACY_OOM; + self->base.ops = &cuda_ops; + self->base.device_type = DAMACY_DEVICE_CUDA; + self->reader = reader; + self->config = *config; + *out = &self->base; + return DAMACY_OK; +} + +void +cuda_resolve_numa(const struct damacy_config* config, struct numa_resolved* out) +{ + *out = (struct numa_resolved){ .node = -1 }; + if (config->tuning.numa_strategy == DAMACY_NUMA_DISABLED || + cuInit(0) != CUDA_SUCCESS) + return; + CUdevice device; + if (config->device >= 0) { + if (cuDeviceGet(&device, config->device) != CUDA_SUCCESS) + return; + } else if (cuCtxGetDevice(&device) != CUDA_SUCCESS) { + return; + } + numa_init( + config->tuning.numa_strategy, config->tuning.numa_node, device, out); +} diff --git a/src/executor/cuda_executor.h b/src/executor/cuda_executor.h new file mode 100644 index 00000000..6e217b42 --- /dev/null +++ b/src/executor/cuda_executor.h @@ -0,0 +1,10 @@ +#pragma once + +#include "damacy_pipeline.h" +#include "numa/numa.h" + +void +cuda_resolve_numa(const struct damacy_config* config, + struct numa_resolved* out); +uint64_t +cuda_executor_set_budget(struct damacy_executor* executor, uint64_t value); diff --git a/src/executor/cuda_geometry.inc b/src/executor/cuda_geometry.inc new file mode 100644 index 00000000..4ef3b40f --- /dev/null +++ b/src/executor/cuda_geometry.inc @@ -0,0 +1,138 @@ +struct resolved_wave_geometry +{ + const struct input_transfer_ops* input; + struct wave_pool_sizing sizing; + struct gpu_budget_breakdown predicted; + uint32_t max_chunks_per_wave; + uint32_t max_substreams_per_chunk; + uint8_t host_buffer_waves; + uint8_t want_gds; +}; + +enum wave_geometry_step +{ + WAVE_GEOMETRY_SIZING, + WAVE_GEOMETRY_PREDICT, +}; + +static enum damacy_status +resolve_wave_geometry(const struct damacy_config* cfg, + uint64_t resolver_budget, + uint64_t runtime_chunk_cap, + struct resolved_wave_geometry* out, + enum wave_geometry_step* failed_step) +{ + *out = (struct resolved_wave_geometry){ 0 }; + out->max_chunks_per_wave = resolve_max_chunks_per_wave(cfg); + out->max_substreams_per_chunk = resolve_max_substreams_per_chunk(cfg); + out->host_buffer_waves = resolve_host_buffer_waves(cfg); + out->want_gds = resolve_enable_gds(cfg); + out->input = + out->want_gds ? input_transfer_gds() : input_transfer_host_staging(); + + const struct input_transfer_resources min_input = + input_transfer_resources(out->input, out->host_buffer_waves, 0); + enum damacy_status s = + wave_pool_resolve_sizing(out->max_chunks_per_wave, + out->max_substreams_per_chunk, + min_input.device_staging_buffers, + resolver_budget, + cfg->tuning.max_read_op_bytes, + runtime_chunk_cap, + cfg->samples_per_batch, + &out->sizing); + if (s != DAMACY_OK) { + if (failed_step) + *failed_step = WAVE_GEOMETRY_SIZING; + return s; + } + + const struct input_transfer_resources input_resources = + input_transfer_resources( + out->input, out->host_buffer_waves, out->sizing.input_staging_per_wave); + s = gpu_budget_predict(cfg, + &input_resources, + out->sizing.dev_decompressed_per_wave, + &out->predicted); + if (s != DAMACY_OK && failed_step) + *failed_step = WAVE_GEOMETRY_PREDICT; + return s; +} + +void +damacy_config_describe(const struct damacy_config* cfg) +{ + if (!cfg) { + log_info("damacy_config_describe: NULL config"); + return; + } + const uint64_t max_gpu = cfg->tuning.max_gpu_memory_bytes; + const uint64_t runtime_chunk_cap = resolve_max_chunk_uncompressed(cfg); + uint64_t pool_reserve = 0; + { + uint64_t pool_bytes = 0; + // resolve_sample_volume_bytes rejects rank=0 / non-positive dims; on + // those, leave pool_reserve at 0 so describe still prints useful + // info for the rest of the geometry. + enum damacy_status pvs = resolve_sample_volume_bytes(cfg, &pool_bytes); + if (pvs == DAMACY_OK) + pool_reserve = 2ull * pool_bytes; + else + log_info( + "damacy_config_describe: resolve_sample_volume_bytes failed (%s); " + "pool_reserve=0", + damacy_status_str(pvs)); + } + const uint64_t resolver_budget = + pool_reserve < max_gpu ? max_gpu - pool_reserve : 0; + log_info("damacy_config_describe: max_gpu_memory_bytes=%llu " + "(pool_reserve=%llu, resolver_budget=%llu, " + "max_chunk_uncompressed_bytes=%llu, samples_per_batch=%u)", + (unsigned long long)max_gpu, + (unsigned long long)pool_reserve, + (unsigned long long)resolver_budget, + (unsigned long long)runtime_chunk_cap, + (unsigned)cfg->samples_per_batch); + + struct resolved_wave_geometry geom = { 0 }; + enum wave_geometry_step failed_step = WAVE_GEOMETRY_SIZING; + enum damacy_status rs = resolve_wave_geometry( + cfg, resolver_budget, runtime_chunk_cap, &geom, &failed_step); + if (rs != DAMACY_OK) { + const char* where = failed_step == WAVE_GEOMETRY_PREDICT + ? "gpu_budget_predict" + : "wave_pool_resolve_sizing"; + log_info( + "damacy_config_describe: %s failed (%s)", where, damacy_status_str(rs)); + return; + } + log_info("damacy_config_describe: input_staging_per_wave=%llu " + "dev_decompressed_per_wave=%llu", + (unsigned long long)geom.sizing.input_staging_per_wave, + (unsigned long long)geom.sizing.dev_decompressed_per_wave); + log_info("damacy_config_describe: dev_compressed=%llu dev_decompressed=%llu " + "blosc1_meta=%llu fanout_soa=%llu " + "nvcomp_temp=%llu batch_metadata=%llu", + (unsigned long long)geom.predicted.dev_compressed, + (unsigned long long)geom.predicted.dev_decompressed, + (unsigned long long)geom.predicted.blosc1_meta, + (unsigned long long)geom.predicted.fanout_soa, + (unsigned long long)geom.predicted.nvcomp_temp, + (unsigned long long)geom.predicted.batch_metadata); + // predicted.total is the *initial* allocation (initial fanout / decoder + // floors); sizing.worst_case_total_bytes is the post-grow worst case + // the resolver pre-reserved against the cap. Their difference is the + // grow-time headroom; the cap minus the worst case is unused slack. + const uint64_t initial_alloc = geom.predicted.total; + const uint64_t worst_case = geom.sizing.worst_case_total_bytes; + const uint64_t reserved_for_grow = + worst_case > initial_alloc ? worst_case - initial_alloc : 0; + const uint64_t slack = max_gpu > worst_case ? max_gpu - worst_case : 0; + log_info("damacy_config_describe: initial_alloc=%llu reserved_for_grow=%llu " + "worst_case_total=%llu slack=%llu cap=%llu", + (unsigned long long)initial_alloc, + (unsigned long long)reserved_for_grow, + (unsigned long long)worst_case, + (unsigned long long)slack, + (unsigned long long)max_gpu); +} diff --git a/src/executor/cuda_stub.c b/src/executor/cuda_stub.c new file mode 100644 index 00000000..514a3985 --- /dev/null +++ b/src/executor/cuda_stub.c @@ -0,0 +1,23 @@ +#include "damacy_pipeline.h" +#include "log/log.h" + +enum damacy_status +damacy_cuda_executor_create(struct damacy_reader* reader, + const struct damacy_cuda_config* config, + struct damacy_executor** out) +{ + (void)reader; + (void)config; + if (!out) + return DAMACY_INVAL; + *out = NULL; + log_error("CUDA support is disabled in this build"); + return DAMACY_CUDA; +} + +void +damacy_config_describe(const struct damacy_config* config) +{ + (void)config; + log_info("CUDA support is disabled in this build"); +} diff --git a/src/executor/dispatch.c b/src/executor/dispatch.c new file mode 100644 index 00000000..52684c0e --- /dev/null +++ b/src/executor/dispatch.c @@ -0,0 +1,140 @@ +#include "executor/dispatch.h" + +#include "executor/coalesce.h" +#include "executor/group_chunks.h" +#include "util/path_intern.h" + +#include +#include + +void +dispatch_scratch_destroy(struct dispatch_scratch* scratch) +{ + free(scratch->indices); + free(scratch->reads); + free(scratch->chunks); + *scratch = (struct dispatch_scratch){ 0 }; +} + +static enum damacy_status +scratch_reserve(struct dispatch_scratch* scratch, uint32_t count) +{ + if (scratch->capacity >= count) + return DAMACY_OK; + dispatch_scratch_destroy(scratch); + scratch->indices = calloc((size_t)count * 4 + 1, sizeof(*scratch->indices)); + scratch->reads = calloc(count, sizeof(*scratch->reads)); + scratch->chunks = calloc(count, sizeof(*scratch->chunks)); + if (!scratch->indices || !scratch->reads || !scratch->chunks) { + dispatch_scratch_destroy(scratch); + return DAMACY_OOM; + } + scratch->capacity = count; + return DAMACY_OK; +} + +enum damacy_status +dispatch_plan_build(const struct prepared_plan* plan, + uint16_t slot, + uint64_t alignment, + uint64_t max_read_bytes, + uint32_t max_chunks_per_wave, + struct dispatch_output* out, + struct dispatch_scratch* scratch) +{ + if (!plan || !out || !scratch || !alignment || !max_chunks_per_wave) + return DAMACY_INVAL; + if (plan->n_uses > out->chunk_plans_cap || plan->n_uses > out->read_ops_cap || + plan->n_regions > out->sample_plans_cap) + return DAMACY_BUDGET; + path_intern_reset(out->paths); + out->n_chunk_plans = out->n_read_ops = out->n_read_op_groups = 0; + out->n_sample_plans = plan->n_regions; + int64_t strides[DAMACY_MAX_RANK + 1]; + strides[plan->output.sample_rank] = 1; + for (int d = plan->output.sample_rank - 1; d >= 0; --d) + strides[d] = strides[d + 1] * plan->output.sample_shape[d]; + for (uint32_t i = 0; i < plan->n_regions; ++i) { + const struct plan_region* region = &plan->regions[i]; + const struct zarr_metadata* meta = &plan->arrays[region->array].metadata; + struct sample_plan* sample = &out->sample_plans[region->sample]; + *sample = + (struct sample_plan){ .batch_pool_slot = slot, + .sample_idx_in_batch = (uint16_t)region->sample, + .rank = meta->rank, + .src_dtype = (uint8_t)meta->dtype, + .sample_dst_off_elems = + (int64_t)region->sample * strides[0], + .chunk_count = 1 }; + memcpy(sample->fill_value, meta->fill_value, sizeof(sample->fill_value)); + int64_t source_stride = 1; + for (int d = meta->rank - 1; d >= 0; --d) { + uint64_t chunk = meta->inner_chunk_shape[d]; + uint64_t begin = (uint64_t)region->source.dims[d].beg / chunk; + uint64_t end = ((uint64_t)region->source.dims[d].end - 1) / chunk + 1; + sample->dims[d] = (struct sample_dim){ + .chunk_shape = (uint32_t)chunk, + .chunk_grid_extent = (uint32_t)(end - begin), + .aabb_lo_relative = + region->source.dims[d].beg - (int64_t)(begin * chunk), + .aabb_extent = region->source.dims[d].end - region->source.dims[d].beg, + .dst_stride = strides[d + 1], + .src_stride = source_stride + }; + sample->chunk_count *= (uint32_t)(end - begin); + source_stride *= (int64_t)chunk; + } + } + const char* source_path = NULL; + const char* dispatch_path = NULL; + for (uint32_t i = 0; i < plan->n_uses; ++i) { + const struct plan_use* use = &plan->uses[i]; + const struct plan_chunk* chunk = &plan->chunks[use->chunk]; + const struct plan_region* region = &plan->regions[use->region]; + const struct zarr_metadata* meta = &plan->arrays[chunk->array].metadata; + struct read_op* read = &out->read_ops[i]; + *read = (struct read_op){ 0 }; + struct chunk_plan* dispatch = &out->chunk_plans[i]; + *dispatch = (struct chunk_plan){ + .read_op_idx = i, + .compressed_nbytes = chunk->encoded_bytes, + .decompressed_nbytes = chunk->decoded_bytes, + .batch_pool_slot = slot, + .sample_idx_in_batch = (uint16_t)region->sample, + .codec_id = chunk->missing ? CODEC_FILL : (uint8_t)meta->inner_codec.id, + .is_fill = chunk->missing + }; + for (uint8_t d = 0; d < meta->rank; ++d) + dispatch->chunk_d[d] = + (uint32_t)(chunk->coordinate[d] - (uint64_t)region->source.dims[d].beg / + meta->inner_chunk_shape[d]); + if (!chunk->missing) { + uint64_t start = chunk->offset / alignment * alignment; + uint64_t end = chunk->offset + chunk->encoded_bytes; + if (end > UINT64_MAX - alignment + 1) + return DAMACY_DECODE; + end = (end + alignment - 1) / alignment * alignment; + if (end - start > UINT32_MAX) + return DAMACY_DECODE; + if (source_path != chunk->path) { + source_path = chunk->path; + dispatch_path = path_intern_acquire(out->paths, source_path); + } + read->shard_path = dispatch_path; + if (!read->shard_path) + return DAMACY_OOM; + read->file_offset = start; + read->nbytes = (uint32_t)(end - start); + dispatch->offset_in_read = (uint32_t)(chunk->offset - start); + } + } + out->n_chunk_plans = out->n_read_ops = plan->n_uses; + enum damacy_status status = scratch_reserve(scratch, plan->n_uses); + if (status != DAMACY_OK) + return status; + status = coalesce_chunks( + out, max_read_bytes, max_chunks_per_wave, scratch->indices, scratch->reads); + if (status != DAMACY_OK) + return status; + return group_chunks_by_read(out, scratch->indices, scratch->chunks); +} diff --git a/src/executor/dispatch.h b/src/executor/dispatch.h new file mode 100644 index 00000000..976cd673 --- /dev/null +++ b/src/executor/dispatch.h @@ -0,0 +1,123 @@ +#pragma once + +#include "planner/plan.h" +#include "zarr/zarr_chunk_layout.h" + +struct path_intern; + +#ifdef __cplusplus +extern "C" +{ +#endif + + struct read_op + { + const char* shard_path; + uint64_t file_offset; + uint64_t host_buf_offset; + uint32_t nbytes; + }; + + struct sample_dim + { + uint32_t chunk_shape; + uint32_t chunk_grid_extent; + int64_t aabb_lo_relative; + int64_t aabb_extent; + int64_t dst_stride; + int64_t src_stride; + }; + + struct sample_plan + { + uint16_t batch_pool_slot; + uint16_t sample_idx_in_batch; + uint8_t rank; + uint8_t src_dtype; + struct sample_dim dims[DAMACY_MAX_RANK]; + int64_t sample_dst_off_elems; + + uint32_t chunk_count; + + uint8_t fill_value[DAMACY_MAX_DTYPE_BYTES]; + + struct chunk_layout layout; + uint8_t layout_probed; + }; + + struct chunk_plan + { + uint32_t read_op_idx; + uint32_t offset_in_read; + uint32_t compressed_nbytes; + uint32_t decompressed_nbytes; + uint64_t host_buf_offset; + uint32_t dev_decompressed_offset; + uint16_t batch_pool_slot; + uint16_t sample_idx_in_batch; + uint8_t codec_id; + uint8_t is_fill; + uint32_t chunk_d[DAMACY_MAX_RANK]; + }; + + struct read_op_group + { + uint32_t read_op_idx; + uint32_t first_chunk; + uint32_t n_chunks; + uint64_t total_decompressed; + }; + + struct read_op_group_iterator + { + const struct read_op_group* groups; + uint32_t n_groups; + uint32_t cursor; + }; + + void read_op_group_iterator_init(struct read_op_group_iterator* it, + const struct read_op_group* groups, + uint32_t n_groups, + uint32_t start_group); + int read_op_group_iterator_next(struct read_op_group_iterator* it, + struct read_op_group* out); + + struct dispatch_output + { + struct read_op* read_ops; + uint32_t read_ops_cap; + uint32_t n_read_ops; + struct chunk_plan* chunk_plans; + uint32_t chunk_plans_cap; + uint32_t n_chunk_plans; + struct sample_plan* sample_plans; + uint32_t sample_plans_cap; + uint32_t n_sample_plans; + struct read_op_group* read_op_groups; + uint32_t read_op_groups_cap; + uint32_t n_read_op_groups; + struct path_intern* paths; + uint32_t n_chunks_to_load; + uint32_t n_loads_issued; + }; + + struct dispatch_scratch + { + uint32_t* indices; + struct read_op* reads; + struct chunk_plan* chunks; + uint32_t capacity; + }; + + enum damacy_status dispatch_plan_build(const struct prepared_plan* plan, + uint16_t slot, + uint64_t alignment, + uint64_t max_read_bytes, + uint32_t max_chunks_per_wave, + struct dispatch_output* out, + struct dispatch_scratch* scratch); + void dispatch_scratch_destroy(struct dispatch_scratch* scratch); + +#ifdef __cplusplus +} +#endif diff --git a/src/planner/group_chunks.c b/src/executor/group_chunks.c similarity index 91% rename from src/planner/group_chunks.c rename to src/executor/group_chunks.c index ef7a9037..9aefa916 100644 --- a/src/planner/group_chunks.c +++ b/src/executor/group_chunks.c @@ -1,9 +1,9 @@ -#include "planner/group_chunks.h" +#include "executor/group_chunks.h" #include enum damacy_status -group_chunks_by_read(struct planner_output* out, +group_chunks_by_read(struct dispatch_output* out, uint32_t* u32_scratch, struct chunk_plan* chunk_plan_scratch) { @@ -52,7 +52,9 @@ group_chunks_by_read(struct planner_output* out, uint32_t r = out->chunk_plans[i].read_op_idx; chunk_plan_scratch[head[r]++] = out->chunk_plans[i]; } - memcpy(out->chunk_plans, chunk_plan_scratch, (size_t)n * sizeof(*chunk_plan_scratch)); + memcpy(out->chunk_plans, + chunk_plan_scratch, + (size_t)n * sizeof(*chunk_plan_scratch)); } for (uint32_t g = 0; g < g_out; ++g) { diff --git a/src/planner/group_chunks.h b/src/executor/group_chunks.h similarity index 86% rename from src/planner/group_chunks.h rename to src/executor/group_chunks.h index 3a3f8b77..b7d6f7b6 100644 --- a/src/planner/group_chunks.h +++ b/src/executor/group_chunks.h @@ -12,8 +12,8 @@ // the result back. #pragma once -#include "damacy.h" // damacy_status -#include "planner/planner.h" // planner_output, chunk_plan +#include "damacy.h" // damacy_status +#include "executor/dispatch.h" // dispatch_output, chunk_plan #include @@ -25,7 +25,7 @@ extern "C" // Scratch: u32_scratch >= n_read_ops+1, chunk_plan_scratch >= n_chunk_plans. // out->read_op_groups must be sized >= out->n_read_ops. enum damacy_status group_chunks_by_read( - struct planner_output* out, + struct dispatch_output* out, uint32_t* u32_scratch, struct chunk_plan* chunk_plan_scratch); diff --git a/src/planner/read_op_sort.c b/src/executor/read_op_sort.c similarity index 96% rename from src/planner/read_op_sort.c rename to src/executor/read_op_sort.c index 688407cf..50d1fe0d 100644 --- a/src/planner/read_op_sort.c +++ b/src/executor/read_op_sort.c @@ -1,6 +1,6 @@ -#include "planner/read_op_sort.h" +#include "executor/read_op_sort.h" -#include "planner/planner.h" // struct read_op +#include "executor/dispatch.h" // struct read_op #include diff --git a/src/planner/read_op_sort.h b/src/executor/read_op_sort.h similarity index 100% rename from src/planner/read_op_sort.h rename to src/executor/read_op_sort.h diff --git a/src/numa/affinity.c b/src/numa/affinity.c new file mode 100644 index 00000000..13aa402e --- /dev/null +++ b/src/numa/affinity.c @@ -0,0 +1,55 @@ +#include "log/log.h" +#include "numa/numa.h" + +#include + +void +numa_scope_enter(const struct numa_resolved* r, struct platform_cpu_mask* saved) +{ + memset(saved, 0, sizeof(*saved)); + if (!r || r->node < 0) + return; + if (platform_thread_affinity_get(saved) != 0) { + log_warn("numa: thread_affinity_get failed; scope is best-effort"); + return; + } + if (platform_thread_affinity_set(&r->cpu_mask) != 0) { + log_warn("numa: thread_affinity_set(node=%d) failed for scope", r->node); + // Clear saved so the matching exit doesn't restore garbage. + memset(saved, 0, sizeof(*saved)); + } +} + +void +numa_scope_exit(const struct platform_cpu_mask* saved) +{ + if (platform_cpu_mask_is_empty(saved)) + return; + if (platform_thread_affinity_set(saved) != 0) + log_warn("numa: thread_affinity_set(restore) failed"); +} + +void +numa_apply_thread_affinity(const struct numa_resolved* r, + const char* thread_label) +{ + if (!r || r->node < 0) + return; + if (platform_thread_affinity_set(&r->cpu_mask) != 0) { + log_warn("numa: thread_affinity_set(node=%d) failed for %s", + r->node, + thread_label ? thread_label : "thread"); + return; + } + struct platform_cpu_mask got; + if (platform_thread_affinity_get(&got) == 0) { + int first = -1, last = -1, count = 0; + if (platform_cpu_mask_describe(&got, &first, &last, &count) == 0) + log_trace("numa: pinned %s to node=%d cpus=[%d..%d] (%d cores)", + thread_label ? thread_label : "thread", + r->node, + first, + last, + count); + } +} diff --git a/src/numa/numa.c b/src/numa/numa.c index ae281be1..daee8af5 100644 --- a/src/numa/numa.c +++ b/src/numa/numa.c @@ -1,4 +1,4 @@ -#include "numa/numa.h" +#include "numa/numa_cuda.h" #include "log/log.h" #include "platform/numa.h" @@ -121,54 +121,3 @@ numa_init(enum damacy_numa_strategy strategy, log_info( "numa: enabled — pinning to host-NUMA node=%d (source=%s)", node, source); } - -void -numa_scope_enter(const struct numa_resolved* r, struct platform_cpu_mask* saved) -{ - memset(saved, 0, sizeof(*saved)); - if (!r || r->node < 0) - return; - if (platform_thread_affinity_get(saved) != 0) { - log_warn("numa: thread_affinity_get failed; scope is best-effort"); - return; - } - if (platform_thread_affinity_set(&r->cpu_mask) != 0) { - log_warn("numa: thread_affinity_set(node=%d) failed for scope", r->node); - // Clear saved so the matching exit doesn't restore garbage. - memset(saved, 0, sizeof(*saved)); - } -} - -void -numa_scope_exit(const struct platform_cpu_mask* saved) -{ - if (platform_cpu_mask_is_empty(saved)) - return; - if (platform_thread_affinity_set(saved) != 0) - log_warn("numa: thread_affinity_set(restore) failed"); -} - -void -numa_apply_thread_affinity(const struct numa_resolved* r, - const char* thread_label) -{ - if (!r || r->node < 0) - return; - if (platform_thread_affinity_set(&r->cpu_mask) != 0) { - log_warn("numa: thread_affinity_set(node=%d) failed for %s", - r->node, - thread_label ? thread_label : "thread"); - return; - } - struct platform_cpu_mask got; - if (platform_thread_affinity_get(&got) == 0) { - int first = -1, last = -1, count = 0; - if (platform_cpu_mask_describe(&got, &first, &last, &count) == 0) - log_trace("numa: pinned %s to node=%d cpus=[%d..%d] (%d cores)", - thread_label ? thread_label : "thread", - r->node, - first, - last, - count); - } -} diff --git a/src/numa/numa.h b/src/numa/numa.h index 2176f811..af804200 100644 --- a/src/numa/numa.h +++ b/src/numa/numa.h @@ -17,8 +17,6 @@ #include "damacy.h" #include "platform/numa.h" -#include - #ifdef __cplusplus extern "C" { @@ -33,13 +31,6 @@ extern "C" struct platform_cpu_mask cpu_mask; }; - // Resolve the GPU's host-NUMA node and populate `out`. Logs once at - // INFO if NUMA is unavailable; that log line is silenced thereafter. - void numa_init(enum damacy_numa_strategy strategy, - int override_node, - CUdevice cu_device, - struct numa_resolved* out); - // Temporarily pin the calling thread to the resolved node's CPU set, // saving the prior mask in `*saved`. Pair with numa_scope_exit. No-op // (and writes an empty `*saved`) when `r->node < 0`. diff --git a/src/numa/numa_cuda.h b/src/numa/numa_cuda.h new file mode 100644 index 00000000..12c48d24 --- /dev/null +++ b/src/numa/numa_cuda.h @@ -0,0 +1,10 @@ +#pragma once + +#include "numa/numa.h" +#include + +void +numa_init(enum damacy_numa_strategy strategy, + int override_node, + CUdevice device, + struct numa_resolved* out); diff --git a/src/pipeline/components.c b/src/pipeline/components.c new file mode 100644 index 00000000..aeffd781 --- /dev/null +++ b/src/pipeline/components.c @@ -0,0 +1,232 @@ +#include "pipeline/components.h" + +#include "damacy_config.h" +#include "log/log.h" + +#include +#include +#include + +enum damacy_status +batch_spec_layout(const struct damacy_batch_spec* spec, + int64_t* shape, + int64_t* strides, + uint64_t* bytes) +{ + if (!spec || !shape || !strides || !bytes || !spec->samples_per_batch || + spec->samples_per_batch > UINT16_MAX || !spec->sample_rank || + spec->sample_rank > DAMACY_MAX_RANK || !damacy_dtype_bpe(spec->dtype)) + return DAMACY_INVAL; + shape[0] = spec->samples_per_batch; + uint64_t elements = spec->samples_per_batch; + for (uint8_t d = 0; d < spec->sample_rank; ++d) { + int64_t extent = spec->sample_shape[d]; + if (extent <= 0 || elements > INT64_MAX / (uint64_t)extent) + return DAMACY_INVAL; + elements *= (uint64_t)extent; + shape[d + 1] = extent; + } + if (elements > SIZE_MAX / damacy_dtype_bpe(spec->dtype)) + return DAMACY_BUDGET; + strides[spec->sample_rank] = 1; + for (int d = spec->sample_rank - 1; d >= 0; --d) + strides[d] = strides[d + 1] * shape[d + 1]; + *bytes = elements * damacy_dtype_bpe(spec->dtype); + return DAMACY_OK; +} + +enum damacy_status +damacy_file_reader_create(uint32_t workers, + uint32_t max_inflight_reads, + struct damacy_reader** out) +{ + if (!out) + return DAMACY_INVAL; + *out = NULL; + if (!workers || workers > DAMACY_MAX_IO_THREADS || !max_inflight_reads) + return DAMACY_INVAL; + struct damacy_reader* reader = calloc(1, sizeof(*reader)); + if (!reader) + return DAMACY_OOM; + reader->max_inflight_reads = max_inflight_reads; + reader->store = store_fs_create( + &(struct store_fs_config){ .root = "", + .nthreads = (int)workers, + .max_inflight_reads = max_inflight_reads }); + if (!reader->store) { + free(reader); + return DAMACY_OOM; + } + *out = reader; + return DAMACY_OK; +} + +void +damacy_reader_destroy(struct damacy_reader* reader) +{ + if (reader) { + store_destroy(reader->store); + free(reader); + } +} + +enum damacy_status +damacy_file_metadata_reader_create(uint32_t concurrency, + const struct damacy_latency_model* latency, + struct damacy_metadata_reader** out) +{ + if (!out) + return DAMACY_INVAL; + *out = NULL; + if (!concurrency || concurrency > DAMACY_MAX_METADATA_IO_CONCURRENCY || + (latency && (!isfinite(latency->lognormal_mu_ln_ns) || + !isfinite(latency->lognormal_sigma_ln_ns) || + latency->lognormal_sigma_ln_ns < 0))) + return DAMACY_INVAL; + struct damacy_metadata_reader* reader = calloc(1, sizeof(*reader)); + if (!reader) + return DAMACY_OOM; + reader->concurrency = concurrency; + if (latency) + reader->latency = *latency; + *out = reader; + return DAMACY_OK; +} + +void +damacy_metadata_reader_destroy(struct damacy_metadata_reader* reader) +{ + if (reader && reader->active) { + log_error("metadata reader is still in use"); + return; + } + free(reader); +} + +enum damacy_status +damacy_zarr_metadata_create(struct damacy_metadata_reader* reader, + const struct damacy_metadata_cache_config* cache, + struct damacy_metadata** out) +{ + if (!out) + return DAMACY_INVAL; + *out = NULL; + if (!reader || !cache || !cache->array_entries || !cache->shard_entries) + return DAMACY_INVAL; + struct damacy_metadata* metadata = calloc(1, sizeof(*metadata)); + if (!metadata) + return DAMACY_OOM; + metadata->reader = reader; + metadata->cache = *cache; + *out = metadata; + return DAMACY_OK; +} + +void +damacy_metadata_destroy(struct damacy_metadata* metadata) +{ + if (metadata && metadata->active) { + log_error("metadata provider is still in use"); + return; + } + free(metadata); +} + +void +damacy_planner_destroy(struct damacy_planner* planner) +{ + if (!planner) + return; + if (planner->active) { + log_error("planner is still in use"); + return; + } + planner->ops->destroy(planner); +} + +void +damacy_executor_destroy(struct damacy_executor* executor) +{ + if (!executor) + return; + if (executor->active) { + log_error("executor is still in use"); + return; + } + executor->ops->destroy(executor); +} + +void +buffer_retain(struct damacy_buffer* buffer) +{ + atomic_fetch_add_explicit(&buffer->references, 1, memory_order_relaxed); +} + +void +buffer_release(struct damacy_buffer* buffer) +{ + if (buffer && atomic_fetch_sub_explicit( + &buffer->references, 1, memory_order_acq_rel) == 1) + buffer->destroy(buffer); +} + +int +buffer_available(const struct damacy_buffer* buffer) +{ + return buffer && + atomic_load_explicit(&buffer->references, memory_order_acquire) == 1; +} + +struct damacy_batch* +batch_create(struct damacy_buffer* buffer, + const struct damacy_batch_spec* output, + uint64_t batch_id) +{ + struct damacy_batch* batch = calloc(1, sizeof(*batch)); + if (!batch) + return NULL; + atomic_init(&batch->references, 1); + batch->buffer = buffer; + buffer_retain(buffer); + batch->info.rank = output->sample_rank + 1; + batch->info.dtype = output->dtype; + batch->info.batch_id = batch_id; + batch->info.shape[0] = output->samples_per_batch; + memcpy(batch->info.shape + 1, + output->sample_shape, + output->sample_rank * sizeof(int64_t)); + return batch; +} + +void +damacy_batch_retain(struct damacy_batch* batch) +{ + if (batch) + atomic_fetch_add_explicit(&batch->references, 1, memory_order_relaxed); +} + +void +damacy_batch_release(struct damacy_batch* batch) +{ + if (batch && atomic_fetch_sub_explicit( + &batch->references, 1, memory_order_acq_rel) == 1) { + buffer_release(batch->buffer); + free(batch); + } +} + +void +damacy_batch_info(const struct damacy_batch* batch, + struct damacy_batch_info* out) +{ + if (!out) + return; + *out = batch ? batch->info : (struct damacy_batch_info){ 0 }; + if (batch) { + out->data = batch->buffer->data; + out->device_ptr = batch->buffer->data; + out->ready_stream = batch->buffer->ready_stream; + out->device_type = batch->buffer->device_type; + out->device_id = batch->buffer->device_id; + } +} diff --git a/src/pipeline/components.h b/src/pipeline/components.h new file mode 100644 index 00000000..5dd3992b --- /dev/null +++ b/src/pipeline/components.h @@ -0,0 +1,113 @@ +#pragma once + +#include "damacy_pipeline.h" +#include "planner/plan.h" +#include "platform/platform.h" +#include "store/store.h" + +#include + +struct damacy_reader +{ + struct store* store; + uint32_t max_inflight_reads; +}; + +struct damacy_metadata_reader +{ + uint32_t concurrency; + struct damacy_latency_model latency; + _Atomic int active; +}; + +struct damacy_metadata +{ + struct damacy_metadata_reader* reader; + struct damacy_metadata_cache_config cache; + _Atomic int active; +}; + +struct damacy_planner_ops +{ + enum damacy_status (*start)(struct damacy_planner*, + const struct damacy_batch_spec*, + const struct damacy_queue_limits*); + struct damacy_push_result (*push)(struct damacy_planner*, + struct damacy_sample_slice); + enum damacy_status (*next)(struct damacy_planner*, struct prepared_plan**); + uint64_t (*pending)(const struct damacy_planner*); + void (*stats)(struct damacy_planner*, struct damacy_stats*); + void (*reset_stats)(struct damacy_planner*); + void (*stop)(struct damacy_planner*); + void (*destroy)(struct damacy_planner*); +}; + +struct damacy_planner +{ + const struct damacy_planner_ops* ops; + _Atomic int active; +}; + +struct damacy_executor_ops +{ + enum damacy_status (*enter_thread)(struct damacy_executor*); + void (*leave_thread)(struct damacy_executor*); + enum damacy_status (*start)(struct damacy_executor*, + const struct damacy_batch_spec*, + struct damacy_stats*); + enum damacy_status (*submit)(struct damacy_executor*, + struct prepared_plan*, + uint64_t); + enum damacy_status (*step)(struct damacy_executor*, int*); + enum damacy_status (*take)(struct damacy_executor*, struct damacy_batch**); + enum damacy_status (*wait_event)(struct damacy_executor*, void*); + int (*busy)(const struct damacy_executor*); + void (*stats)(struct damacy_executor*, struct damacy_stats*); + void (*stop)(struct damacy_executor*); + void (*destroy)(struct damacy_executor*); +}; + +struct damacy_executor +{ + const struct damacy_executor_ops* ops; + _Atomic int active; + int device_type; + int device_id; +}; + +struct damacy_buffer +{ + _Atomic uint32_t references; + void* data; + uint64_t nbytes; + void* ready_stream; + int device_type; + int device_id; + void (*destroy)(struct damacy_buffer*); + enum damacy_status (*wait_event)(struct damacy_buffer*, void*); + void* context; +}; + +struct damacy_batch +{ + _Atomic uint32_t references; + struct damacy_buffer* buffer; + struct damacy_batch_info info; + const struct damacy* owner; +}; + +enum damacy_status +batch_spec_layout(const struct damacy_batch_spec* spec, + int64_t* shape, + int64_t* strides, + uint64_t* bytes); +struct damacy_batch* +batch_create(struct damacy_buffer* buffer, + const struct damacy_batch_spec* output, + uint64_t batch_id); +void +buffer_retain(struct damacy_buffer* buffer); +void +buffer_release(struct damacy_buffer* buffer); +int +buffer_available(const struct damacy_buffer* buffer); diff --git a/src/pipeline/zarr_planner.c b/src/pipeline/zarr_planner.c new file mode 100644 index 00000000..edf45c86 --- /dev/null +++ b/src/pipeline/zarr_planner.c @@ -0,0 +1,320 @@ +#include "pipeline/components.h" + +#include "log/log.h" +#include "lookahead/lookahead.h" +#include "planner/plan_builder.h" +#include "prefetch/array_meta.h" +#include "prefetch/prefetcher.h" +#include "prefetch/shard_index.h" +#include "store/metadata_store_async.h" + +#include + +struct zarr_planner +{ + struct damacy_planner base; + struct damacy_metadata* metadata; + struct damacy_plan_limits limits; + struct damacy_batch_spec output; + struct damacy_queue_limits queues; + struct metadata_store_async* reader; + struct array_meta_async_fetcher array_fetcher; + struct shard_index_async_fetcher shard_fetcher; + struct prefetch_cache* arrays; + struct prefetch_cache* shards; + struct damacy_lookahead lookahead; + struct prefetcher* prefetcher; + struct planner_sample* samples; + uint32_t staged; + uint64_t pushed; + uint64_t planned; + uint64_t watermark; + int bound; +}; + +static void +clear_samples(struct zarr_planner* self) +{ + for (uint32_t i = 0; i < self->staged; ++i) { + free((char*)self->samples[i].uri); + free(self->samples[i].h_shards); + self->samples[i] = (struct planner_sample){ 0 }; + } + self->staged = 0; +} + +static void +zarr_stop(struct damacy_planner* base) +{ + struct zarr_planner* self = (void*)base; + prefetcher_stop(self->prefetcher); + metadata_store_async_destroy(self->reader); + self->reader = NULL; + prefetcher_destroy(self->prefetcher); + self->prefetcher = NULL; + prefetch_cache_destroy(self->shards); + self->shards = NULL; + prefetch_cache_destroy(self->arrays); + self->arrays = NULL; + lookahead_destroy(&self->lookahead); + clear_samples(self); + free(self->samples); + self->samples = NULL; + if (self->bound) { + self->metadata->active = 0; + self->metadata->reader->active = 0; + self->bound = 0; + } +} + +static enum damacy_status +zarr_start(struct damacy_planner* base, + const struct damacy_batch_spec* output, + const struct damacy_queue_limits* queues) +{ + struct zarr_planner* self = (void*)base; + struct damacy_metadata* metadata = self->metadata; + if (metadata->active || metadata->reader->active) + return DAMACY_INVAL; + uint64_t floor = + (uint64_t)queues->lookahead_samples + output->samples_per_batch; + if (metadata->cache.array_entries < floor || + metadata->cache.shard_entries < + floor * self->limits.max_shards_per_sample) { + log_error( + "metadata cache requires array_entries >= %llu and shard_entries >= %llu", + (unsigned long long)floor, + (unsigned long long)(floor * self->limits.max_shards_per_sample)); + return DAMACY_INVAL; + } + self->output = *output; + self->queues = *queues; + self->pushed = self->planned = self->watermark = 0; + int expected = 0; + if (!atomic_compare_exchange_strong(&metadata->active, &expected, 1)) + return DAMACY_INVAL; + expected = 0; + if (!atomic_compare_exchange_strong( + &metadata->reader->active, &expected, 1)) { + metadata->active = 0; + return DAMACY_INVAL; + } + self->bound = 1; + self->reader = metadata_store_async_create( + (int)metadata->reader->concurrency, NULL, &metadata->reader->latency); + if (!self->reader) + goto Fail; + array_meta_async_fetcher_init(&self->array_fetcher, self->reader); + self->arrays = prefetch_cache_create(&(struct prefetch_cache_config){ + .capacity = metadata->cache.array_entries, + .max_probe = 16, + .knob_name = "array_entries", + .ops = &array_meta_ops, + .async_fetcher = &self->array_fetcher.base }); + if (!self->arrays) + goto Fail; + shard_index_async_fetcher_init( + &self->shard_fetcher, self->reader, self->arrays); + self->shards = prefetch_cache_create(&(struct prefetch_cache_config){ + .capacity = metadata->cache.shard_entries, + .max_probe = 16, + .knob_name = "shard_entries", + .ops = &shard_index_ops, + .async_fetcher = &self->shard_fetcher.base }); + if (!self->shards || + lookahead_init(&self->lookahead, queues->lookahead_samples)) + goto Fail; + self->samples = calloc(output->samples_per_batch, sizeof(*self->samples)); + if (!self->samples) + goto Fail; + self->prefetcher = prefetcher_create(&(struct prefetcher_config){ + .lookahead = &self->lookahead, + .array_meta_cache = self->arrays, + .shard_index_cache = self->shards, + .capacity = queues->lookahead_samples, + .owner_capacity = queues->lookahead_samples + 4, + .max_shards_per_sample = self->limits.max_shards_per_sample }); + if (!self->prefetcher || prefetcher_start(self->prefetcher)) + goto Fail; + return DAMACY_OK; +Fail: + zarr_stop(base); + return DAMACY_OOM; +} + +static struct damacy_push_result +zarr_push(struct damacy_planner* base, struct damacy_sample_slice samples) +{ + struct zarr_planner* self = (void*)base; + struct damacy_push_result result = { .unconsumed = samples, + .status = DAMACY_OK }; + for (const struct damacy_sample* sample = samples.beg; sample != samples.end; + ++sample) { + result.unconsumed.beg = sample; + if (prefetcher_unconsumed_count(self->prefetcher, self->pushed) >= + self->queues.lookahead_samples) { + result.status = DAMACY_AGAIN; + return result; + } + if (!sample->uri) { + result.status = DAMACY_INVAL; + return result; + } + if (sample->aabb.rank != self->output.sample_rank) { + result.status = DAMACY_RANK; + return result; + } + for (uint8_t d = 0; d < sample->aabb.rank; ++d) { + int64_t lo = sample->aabb.dims[d].beg; + int64_t hi = sample->aabb.dims[d].end; + if (lo < 0 || hi <= lo || hi - lo != self->output.sample_shape[d]) { + result.status = DAMACY_INVAL; + return result; + } + } + if (lookahead_push_with_sample_seq( + &self->lookahead, sample, self->pushed)) { + result.status = DAMACY_OOM; + return result; + } + ++self->pushed; + } + result.unconsumed.beg = samples.end; + return result; +} + +static enum damacy_status +zarr_next(struct damacy_planner* base, struct prepared_plan** out) +{ + struct zarr_planner* self = (void*)base; + *out = NULL; + while (self->staged < self->output.samples_per_batch) { + struct prefetcher_ready ready = { 0 }; + if (!prefetcher_pop_ready(self->prefetcher, &ready)) + return DAMACY_AGAIN; + if (ready.result == PREFETCHER_RESULT_ERROR) { + enum damacy_status status = + ready.err_code ? (enum damacy_status)ready.err_code : DAMACY_INVAL; + prefetcher_ready_free(&ready); + return status; + } + self->samples[self->staged++] = + (struct planner_sample){ .uri = ready.uri, + .aabb = ready.aabb, + .h_meta = ready.h_meta, + .h_shards = ready.h_shards, + .n_shards = ready.n_shards }; + self->watermark = ready.sample_seq + 1; + } + enum damacy_status status = prepared_plan_build(self->arrays, + self->shards, + self->samples, + self->staged, + &self->output, + &self->limits, + out); + if (status == DAMACY_AGAIN) + return status; + self->planned += self->staged; + clear_samples(self); + prefetcher_advance_watermark(self->prefetcher, self->watermark); + return status; +} + +static uint64_t +zarr_pending(const struct damacy_planner* base) +{ + const struct zarr_planner* self = (const void*)base; + return self->pushed - self->planned; +} + +static void +zarr_stats(struct damacy_planner* base, struct damacy_stats* out) +{ + struct zarr_planner* self = (void*)base; + struct prefetch_cache_stats cache; + if (self->arrays) { + prefetch_cache_stats_get(self->arrays, &cache); + out->array_meta.hits = cache.counters.hits; + out->array_meta.misses = cache.counters.misses; + } + if (self->shards) { + prefetch_cache_stats_get(self->shards, &cache); + out->shard_index.hits = cache.counters.hits; + out->shard_index.misses = cache.counters.misses; + } + if (!self->reader) + return; + struct metadata_store_async_latency_stats latency; + metadata_store_async_latency_stats_get(self->reader, &latency); + out->metadata_latency.ops = latency.ops; + out->metadata_latency.stat_ops = latency.stat_ops; + out->metadata_latency.submit_ops = latency.submit_ops; + out->metadata_latency.active = latency.active; + out->metadata_latency.max_active = latency.max_active; + out->metadata_latency.total_sleep_ns = latency.total_sleep_ns; + out->metadata_latency.max_sleep_ns = latency.max_sleep_ns; + struct metadata_store_async_backend_stats backend; + metadata_store_async_backend_stats_get(self->reader, &backend); + out->metadata_backend.read_jobs = backend.read_jobs; + out->metadata_backend.read_active = backend.read_active; + out->metadata_backend.read_max_active = backend.read_max_active; + struct metadata_store_async_op_latency_stats operations; + metadata_store_async_op_latency_stats_get(self->reader, &operations); + for (unsigned i = 0; i < DAMACY_METADATA_OP_LATENCY_NKINDS; ++i) { + out->metadata_op_latency[i].count = operations.kinds[i].count; + out->metadata_op_latency[i].sum_ns = operations.kinds[i].sum_ns; + out->metadata_op_latency[i].max_ns = operations.kinds[i].max_ns; + for (unsigned j = 0; j < DAMACY_METADATA_OP_LATENCY_NBUCKETS; ++j) + out->metadata_op_latency[i].buckets[j] = operations.kinds[i].buckets[j]; + } +} + +static void +zarr_reset_stats(struct damacy_planner* base) +{ + struct zarr_planner* self = (void*)base; + metadata_store_async_latency_stats_reset(self->reader); + metadata_store_async_backend_stats_reset(self->reader); + metadata_store_async_op_latency_stats_reset(self->reader); +} + +static void +zarr_destroy(struct damacy_planner* base) +{ + zarr_stop(base); + free(base); +} + +static const struct damacy_planner_ops zarr_ops = { .start = zarr_start, + .push = zarr_push, + .next = zarr_next, + .pending = zarr_pending, + .stats = zarr_stats, + .reset_stats = + zarr_reset_stats, + .stop = zarr_stop, + .destroy = zarr_destroy }; + +enum damacy_status +damacy_chunk_planner_create(struct damacy_metadata* metadata, + const struct damacy_plan_limits* limits, + struct damacy_planner** out) +{ + if (!out) + return DAMACY_INVAL; + *out = NULL; + if (!metadata || !limits || !limits->max_chunks || + limits->max_chunks > DAMACY_MAX_CHUNKS_PER_BATCH || + !limits->max_chunk_bytes || !limits->max_shards_per_sample || + !limits->max_plan_bytes) + return DAMACY_INVAL; + struct zarr_planner* self = calloc(1, sizeof(*self)); + if (!self) + return DAMACY_OOM; + self->base.ops = &zarr_ops; + self->metadata = metadata; + self->limits = *limits; + *out = &self->base; + return DAMACY_OK; +} diff --git a/src/planner/plan.c b/src/planner/plan.c new file mode 100644 index 00000000..3e07faaa --- /dev/null +++ b/src/planner/plan.c @@ -0,0 +1,12 @@ +#include "planner/plan.h" + +#include + +void +prepared_plan_destroy(struct prepared_plan* plan) +{ + if (!plan) + return; + free(plan->storage); + free(plan); +} diff --git a/src/planner/plan.h b/src/planner/plan.h new file mode 100644 index 00000000..2b41b2b0 --- /dev/null +++ b/src/planner/plan.h @@ -0,0 +1,60 @@ +#pragma once + +#include "damacy_pipeline.h" +#include "zarr/zarr_metadata.h" + +struct plan_array +{ + const char* uri; + struct zarr_metadata metadata; +}; + +struct plan_chunk +{ + const char* path; + uint64_t offset; + uint64_t coordinate[DAMACY_MAX_RANK]; + uint32_t array; + uint32_t encoded_bytes; + uint32_t decoded_bytes; + uint32_t first_use; + uint8_t missing; +}; + +enum plan_operation +{ + PLAN_COPY, +}; + +struct plan_region +{ + enum plan_operation operation; + uint32_t array; + uint32_t sample; + struct damacy_aabb source; +}; + +struct plan_use +{ + uint32_t chunk; + uint32_t region; + uint32_t next; +}; + +struct prepared_plan +{ + struct damacy_batch_spec output; + struct plan_array* arrays; + struct plan_chunk* chunks; + struct plan_region* regions; + struct plan_use* uses; + uint32_t n_arrays; + uint32_t n_chunks; + uint32_t n_regions; + uint32_t n_uses; + uint64_t allocated_bytes; + void* storage; +}; + +void +prepared_plan_destroy(struct prepared_plan* plan); diff --git a/src/planner/plan_builder.c b/src/planner/plan_builder.c new file mode 100644 index 00000000..7d7d8f35 --- /dev/null +++ b/src/planner/plan_builder.c @@ -0,0 +1,316 @@ +#include "planner/plan_builder.h" + +#include "damacy_config.h" +#include "prefetch/prefetch_cache.h" +#include "prefetch/shard_index.h" +#include "util/hash.h" +#include "util/strbuf.h" +#include "zarr/sample_shard_iterator.h" + +#include +#include + +static enum damacy_status +array_metadata(struct prefetch_cache* cache, + struct prefetch_handle handle, + const struct zarr_metadata** out) +{ + const void* value = NULL; + int error = 0; + enum prefetch_state state = + prefetch_cache_query(cache, handle, &value, &error); + if (state == PREFETCH_STATE_PENDING) + return DAMACY_AGAIN; + if (state == PREFETCH_STATE_ERROR) + return error ? (enum damacy_status)error : DAMACY_INVAL; + *out = value; + return value ? DAMACY_OK : DAMACY_INVAL; +} + +static enum damacy_status +sample_geometry(const struct planner_sample* sample, + const struct zarr_metadata* meta, + const struct damacy_batch_spec* output, + const struct damacy_plan_limits* limits, + uint64_t* begin, + uint64_t* end, + uint32_t* count, + uint32_t* bytes) +{ + if (sample->aabb.rank != meta->rank || meta->rank != output->sample_rank) + return DAMACY_RANK; + if (!cast_path_supported(output->dtype, meta->dtype)) + return DAMACY_DTYPE; + uint64_t total = 1; + uint64_t size = dtype_bpe(meta->dtype); + for (uint8_t d = 0; d < meta->rank; ++d) { + int64_t lo = sample->aabb.dims[d].beg; + int64_t hi = sample->aabb.dims[d].end; + uint64_t chunk = meta->inner_chunk_shape[d]; + if (lo < 0 || hi <= lo || (uint64_t)hi > meta->shape[d] || !chunk) + return DAMACY_INVAL; + if (hi - lo != output->sample_shape[d]) + return DAMACY_INVAL; + begin[d] = (uint64_t)lo / chunk; + end[d] = ((uint64_t)hi - 1) / chunk + 1; + if (total > limits->max_chunks / (end[d] - begin[d])) + return DAMACY_BUDGET; + total *= end[d] - begin[d]; + if (size > limits->max_chunk_bytes / chunk) + return DAMACY_BUDGET; + size *= chunk; + } + *count = (uint32_t)total; + *bytes = (uint32_t)size; + return DAMACY_OK; +} + +static uint64_t +chunk_hash(uint32_t array, const uint64_t* coordinate, uint8_t rank) +{ + uint64_t hash = array; + for (uint8_t d = 0; d < rank; ++d) + hash = hash_combine(hash, coordinate[d]); + return hash; +} + +static char* +copy_path(char** cursor, const char* path) +{ + size_t length = strlen(path) + 1; + char* result = *cursor; + memcpy(result, path, length); + *cursor += length; + return result; +} + +enum damacy_status +prepared_plan_build(struct prefetch_cache* arrays, + struct prefetch_cache* shards, + const struct planner_sample* samples, + uint32_t n_samples, + const struct damacy_batch_spec* output, + const struct damacy_plan_limits* limits, + struct prepared_plan** out) +{ + if (!out) + return DAMACY_INVAL; + *out = NULL; + if (!arrays || !shards || !samples || !output || !limits || !n_samples || + !limits->max_chunks || !limits->max_chunk_bytes || + output->sample_rank == 0 || output->sample_rank > DAMACY_MAX_RANK) + return DAMACY_INVAL; + enum damacy_status status = DAMACY_OK; + struct prepared_plan* plan = NULL; + struct strbuf path = { 0 }; + uint64_t capacity = 0; + uint64_t path_bytes = 0; + for (uint32_t i = 0; i < n_samples; ++i) { + const struct zarr_metadata* meta = NULL; + if (!samples[i].uri || samples[i].n_shards > limits->max_shards_per_sample) + return DAMACY_INVAL; + status = array_metadata(arrays, samples[i].h_meta, &meta); + if (status != DAMACY_OK) + return status; + uint64_t begin[DAMACY_MAX_RANK], end[DAMACY_MAX_RANK]; + uint32_t count, bytes; + status = sample_geometry( + &samples[i], meta, output, limits, begin, end, &count, &bytes); + if (status != DAMACY_OK) + return status; + capacity += count; + if (capacity > limits->max_chunks) + return DAMACY_BUDGET; + uint64_t length = strlen(samples[i].uri) + 1; + uint64_t per_path = length + 3 + (uint64_t)meta->rank * 21; + if (per_path > limits->max_plan_bytes || + samples[i].n_shards > limits->max_plan_bytes / per_path) + return DAMACY_BUDGET; + uint64_t addition = length + per_path * samples[i].n_shards; + if (addition > limits->max_plan_bytes || + path_bytes > limits->max_plan_bytes - addition) + return DAMACY_BUDGET; + path_bytes += addition; + } + uint64_t buckets = 1; + while (buckets < 2 * capacity) + buckets *= 2; + uint64_t storage_bytes = + (uint64_t)n_samples * + (sizeof(struct plan_array) + sizeof(struct plan_region)) + + capacity * (sizeof(struct plan_chunk) + sizeof(struct plan_use)) + + buckets * sizeof(uint32_t) + path_bytes; + if (storage_bytes > SIZE_MAX || storage_bytes > limits->max_plan_bytes || + sizeof(*plan) > limits->max_plan_bytes - storage_bytes) + return DAMACY_BUDGET; + plan = calloc(1, sizeof(*plan)); + if (!plan) + return DAMACY_OOM; + plan->storage = calloc(1, (size_t)storage_bytes); + if (!plan->storage) { + status = DAMACY_OOM; + goto Done; + } + plan->allocated_bytes = sizeof(*plan) + storage_bytes; + plan->output = *output; + plan->arrays = plan->storage; + plan->regions = (void*)(plan->arrays + n_samples); + plan->chunks = (void*)(plan->regions + n_samples); + plan->uses = (void*)(plan->chunks + capacity); + uint32_t* table = (void*)(plan->uses + capacity); + char* paths = (void*)(table + buckets); + for (uint32_t i = 0; i < n_samples; ++i) { + const struct planner_sample* sample = &samples[i]; + const struct zarr_metadata* meta = NULL; + status = array_metadata(arrays, sample->h_meta, &meta); + if (status != DAMACY_OK) + goto Done; + uint32_t array = 0; + while (array < plan->n_arrays && + strcmp(plan->arrays[array].uri, sample->uri)) + ++array; + if (array == plan->n_arrays) { + plan->arrays[array].uri = copy_path(&paths, sample->uri); + plan->arrays[array].metadata = *meta; + ++plan->n_arrays; + } + plan->regions[i] = (struct plan_region){ .operation = PLAN_COPY, + .array = array, + .sample = i, + .source = sample->aabb }; + ++plan->n_regions; + uint64_t begin[DAMACY_MAX_RANK], end[DAMACY_MAX_RANK]; + uint64_t per_shard[DAMACY_MAX_RANK]; + uint32_t count, decoded_bytes; + status = sample_geometry( + sample, meta, output, limits, begin, end, &count, &decoded_bytes); + if (status != DAMACY_OK) + goto Done; + if (zarr_metadata_inner_per_shard(meta, per_shard, NULL)) { + status = DAMACY_DECODE; + goto Done; + } + struct sample_shard_iterator iterator; + if (sample_shard_iterator_init(&iterator, meta, &sample->aabb)) { + status = DAMACY_INVAL; + goto Done; + } + uint64_t shard_coord[DAMACY_MAX_RANK]; + uint32_t shard_index = 0; + while (sample_shard_iterator_next(&iterator, shard_coord)) { + if (shard_index >= sample->n_shards) { + status = DAMACY_INVAL; + goto Done; + } + const void* value = NULL; + int error = 0; + enum prefetch_state state = prefetch_cache_query( + shards, sample->h_shards[shard_index++], &value, &error); + int missing = state == PREFETCH_STATE_ERROR && error == DAMACY_NOTFOUND; + if (!missing && (state != PREFETCH_STATE_READY || !value)) { + status = state == PREFETCH_STATE_PENDING ? DAMACY_AGAIN + : error ? (enum damacy_status)error + : DAMACY_DECODE; + goto Done; + } + const struct shard_index_value* index = value; + const char* shard_path = NULL; + if (!missing) { + if (zarr_shard_path_build( + &path, sample->uri, shard_coord, meta->rank)) { + status = DAMACY_OOM; + goto Done; + } + shard_path = copy_path(&paths, strbuf_cstr(&path)); + } + uint64_t lo[DAMACY_MAX_RANK], hi[DAMACY_MAX_RANK]; + uint64_t coordinate[DAMACY_MAX_RANK] = { 0 }; + for (uint8_t d = 0; d < meta->rank; ++d) { + uint64_t shard_lo = shard_coord[d] * per_shard[d]; + uint64_t shard_hi = shard_lo + per_shard[d]; + lo[d] = begin[d] > shard_lo ? begin[d] : shard_lo; + hi[d] = end[d] < shard_hi ? end[d] : shard_hi; + coordinate[d] = lo[d]; + } + for (;;) { + uint64_t entry_index = 0; + for (uint8_t d = 0; d < meta->rank; ++d) + entry_index = + entry_index * per_shard[d] + coordinate[d] % per_shard[d]; + struct plan_chunk chunk = { .array = array, + .path = shard_path, + .decoded_bytes = decoded_bytes, + .first_use = UINT32_MAX, + .missing = (uint8_t)missing }; + memcpy(chunk.coordinate, coordinate, sizeof(coordinate)); + if (!missing) { + if (entry_index >= index->n_entries) { + status = DAMACY_DECODE; + goto Done; + } + const struct zarr_shard_entry* entry = &index->entries[entry_index]; + int empty_offset = entry->offset == ZARR_SHARD_EMPTY_OFFSET; + int empty_size = entry->nbytes == ZARR_SHARD_EMPTY_NBYTES; + if (empty_offset != empty_size || + (!empty_size && (!entry->nbytes || entry->nbytes > UINT32_MAX || + entry->offset > UINT64_MAX - entry->nbytes))) { + status = DAMACY_DECODE; + goto Done; + } + chunk.missing = (uint8_t)empty_size; + if (!empty_size) { + chunk.offset = entry->offset; + chunk.encoded_bytes = (uint32_t)entry->nbytes; + } + } + uint64_t bucket = + chunk_hash(array, coordinate, meta->rank) & (buckets - 1); + while (table[bucket]) { + const struct plan_chunk* found = &plan->chunks[table[bucket] - 1]; + if (found->array == array && + !memcmp(found->coordinate, coordinate, sizeof(coordinate))) + break; + bucket = (bucket + 1) & (buckets - 1); + } + uint32_t chunk_index; + if (!table[bucket]) { + chunk_index = plan->n_chunks++; + plan->chunks[chunk_index] = chunk; + table[bucket] = chunk_index + 1; + } else { + chunk_index = table[bucket] - 1; + } + if (plan->n_uses >= capacity) { + status = DAMACY_BUDGET; + goto Done; + } + struct plan_chunk* stored = &plan->chunks[chunk_index]; + plan->uses[plan->n_uses] = (struct plan_use){ + .chunk = chunk_index, .region = i, .next = stored->first_use + }; + stored->first_use = plan->n_uses++; + int finished = 1; + for (int d = meta->rank - 1; d >= 0; --d) { + if (++coordinate[d] < hi[d]) { + finished = 0; + break; + } + coordinate[d] = lo[d]; + } + if (finished) + break; + } + } + if (shard_index != sample->n_shards) { + status = DAMACY_INVAL; + goto Done; + } + } +Done: + strbuf_free(&path); + if (status != DAMACY_OK) + prepared_plan_destroy(plan); + else + *out = plan; + return status; +} diff --git a/src/planner/plan_builder.h b/src/planner/plan_builder.h new file mode 100644 index 00000000..f7da85d0 --- /dev/null +++ b/src/planner/plan_builder.h @@ -0,0 +1,25 @@ +#pragma once + +#include "planner/plan.h" +#include "prefetch/prefetch_handle.h" + +struct planner_sample +{ + const char* uri; + struct damacy_aabb aabb; + struct prefetch_handle h_meta; + struct prefetch_handle* h_shards; + uint32_t n_shards; + struct prefetch_handle h_layout; +}; + +struct prefetch_cache; + +enum damacy_status +prepared_plan_build(struct prefetch_cache* arrays, + struct prefetch_cache* shards, + const struct planner_sample* samples, + uint32_t n_samples, + const struct damacy_batch_spec* output, + const struct damacy_plan_limits* limits, + struct prepared_plan** out); diff --git a/src/planner/planner.c b/src/planner/planner.c index 83e0b666..aee32e36 100644 --- a/src/planner/planner.c +++ b/src/planner/planner.c @@ -1,671 +1,158 @@ #include "planner/planner.h" -#include "damacy_config.h" -#include "dtype/dtype.h" -#include "log/log.h" -#include "planner/coalesce.h" -#include "planner/group_chunks.h" #include "prefetch/prefetch_cache.h" -#include "prefetch/shard_index.h" -#include "util/path_intern.h" -#include "util/prelude.h" -#include "util/strbuf.h" -#include "zarr/sample_shard_iterator.h" -#include "zarr/zarr_metadata.h" -#include "zarr/zarr_shard_index.h" #include #include struct planner { - struct planner_config cfg; - struct strbuf path_sb; - uint32_t* scratch_u32; - uint32_t scratch_u32_cap; - struct read_op* scratch_ops; - uint32_t scratch_ops_cap; - struct chunk_plan* scratch_chunk_plans; - uint32_t scratch_chunk_plans_cap; + struct planner_config config; + struct dispatch_scratch scratch; }; -// `need` = pre-coalesce n_read_ops (== n_chunk_plans). The uint32 -// buffer is sized to 4*need; post-coalesce n_read_ops <= need so -// group_chunks_by_read's n_read_ops+1 slots fit in the same buffer. -static enum damacy_status -planner_ensure_scratch(struct planner* self, uint32_t need) -{ - if (need > self->scratch_u32_cap) { - free(self->scratch_u32); - self->scratch_u32 = NULL; - self->scratch_u32_cap = 0; - uint32_t* mem = (uint32_t*)malloc((size_t)need * 4u * sizeof(uint32_t)); - if (!mem) - return DAMACY_OOM; - self->scratch_u32 = mem; - self->scratch_u32_cap = need; - } - if (need > self->scratch_ops_cap) { - free(self->scratch_ops); - self->scratch_ops = NULL; - self->scratch_ops_cap = 0; - struct read_op* mem = - (struct read_op*)malloc((size_t)need * sizeof(struct read_op)); - if (!mem) - return DAMACY_OOM; - self->scratch_ops = mem; - self->scratch_ops_cap = need; - } - if (need > self->scratch_chunk_plans_cap) { - free(self->scratch_chunk_plans); - self->scratch_chunk_plans = NULL; - self->scratch_chunk_plans_cap = 0; - struct chunk_plan* mem = - (struct chunk_plan*)malloc((size_t)need * sizeof(struct chunk_plan)); - if (!mem) - return DAMACY_OOM; - self->scratch_chunk_plans = mem; - self->scratch_chunk_plans_cap = need; - } - return DAMACY_OK; -} - -// --- math helpers -------------------------------------------------------- - -// chunk_beg/chunk_end are inner-chunk-grid units (half-open). -static int -chunk_range(const struct damacy_aabb* aabb, - const struct zarr_metadata* meta, - uint64_t* chunk_beg, - uint64_t* chunk_end) -{ - if (aabb->rank != meta->rank) - return 1; - for (uint8_t d = 0; d < meta->rank; ++d) { - int64_t beg = aabb->dims[d].beg; - int64_t end = aabb->dims[d].end; - if (beg < 0 || end <= beg) - return 1; - if ((uint64_t)end > meta->shape[d]) - return 1; - uint64_t chunk_extent = meta->inner_chunk_shape[d]; - if (chunk_extent == 0) - return 1; - chunk_beg[d] = (uint64_t)beg / chunk_extent; - chunk_end[d] = ((uint64_t)end - 1) / chunk_extent + 1; - } - return 0; -} - -// Row-major linear index into a multi-d grid given per-dim extents. -static uint64_t -row_major_linear(const uint64_t* coord, const uint64_t* extents, uint8_t rank) -{ - uint64_t linear_idx = 0; - for (uint8_t d = 0; d < rank; ++d) - linear_idx = linear_idx * extents[d] + coord[d]; - return linear_idx; -} - -// Element strides (row-major) into a row-major chunk of shape `shape`. -static void -row_major_strides(const uint64_t* shape, uint8_t rank, int64_t* out_strides) -{ - if (rank == 0) - return; - out_strides[rank - 1] = 1; - for (int d = (int)rank - 2; d >= 0; --d) - out_strides[d] = out_strides[d + 1] * (int64_t)shape[d + 1]; -} - -// Bytes per inner chunk (uncompressed). Returns 0 on overflow during -// the product or if the result exceeds DAMACY_MAX_CHUNK_BYTES. -static uint64_t -inner_chunk_bytes(const struct zarr_metadata* meta) -{ - uint64_t total_bytes = dtype_bpe(meta->dtype); - for (uint8_t d = 0; d < meta->rank; ++d) { - uint64_t chunk_extent = meta->inner_chunk_shape[d]; - if (chunk_extent != 0 && total_bytes > UINT64_MAX / chunk_extent) - return 0; - total_bytes *= chunk_extent; - } - if (total_bytes > DAMACY_MAX_CHUNK_BYTES) - return 0; - return total_bytes; -} - -static uint64_t -align_down_u64(uint64_t value, uint64_t alignment) -{ - return (value / alignment) * alignment; -} - -static uint64_t -align_up_u64(uint64_t value, uint64_t alignment) -{ - return ((value + alignment - 1) / alignment) * alignment; -} - -// --- public API ---------------------------------------------------------- - enum damacy_status -planner_create(const struct planner_config* cfg, struct planner** out) +planner_create(const struct planner_config* config, struct planner** out) { - struct planner* self = NULL; - enum damacy_status status = DAMACY_INVAL; - - CHECK_SILENT(Error, out); + if (!out) + return DAMACY_INVAL; *out = NULL; - CHECK_SILENT(Error, cfg); - CHECK_SILENT(Error, cfg->array_meta_cache); - CHECK_SILENT(Error, cfg->chunk_layout_cache); - CHECK_SILENT(Error, cfg->shard_index_cache); - CHECK_SILENT(Error, cfg->page_alignment > 0); - CHECK_SILENT(Error, cfg->max_chunks_per_wave > 0); - CHECK_SILENT(Error, cfg->max_substreams_per_chunk > 0); - - status = DAMACY_OOM; - self = (struct planner*)calloc(1, sizeof(*self)); - CHECK(Error, self); - - // Designated init zeroes any fields not explicitly named, including - // path_sb (zero-init is documented as safe for struct strbuf). - *self = (struct planner){ .cfg = *cfg }; + if (!config || !config->array_meta_cache || !config->shard_index_cache || + !config->chunk_layout_cache || !config->page_alignment || + !config->max_chunks_per_wave || !config->max_substreams_per_chunk) + return DAMACY_INVAL; + struct planner* self = calloc(1, sizeof(*self)); + if (!self) + return DAMACY_OOM; + self->config = *config; *out = self; return DAMACY_OK; - -Error: - planner_destroy(self); - return status; } void planner_destroy(struct planner* self) { - if (!self) - return; - strbuf_free(&self->path_sb); - free(self->scratch_u32); - free(self->scratch_ops); - free(self->scratch_chunk_plans); - free(self); -} - -// Per-sample/per-shard invariants for a run of emit_chunk calls. The -// per-sample fields stay constant while iterating chunks; per-shard -// fields are refreshed when the iterator crosses into a new shard. -// -// shard_missing == 1 short-circuits the per-shard entry lookup and emits -// fill chunks for every chunk inside the shard (shard file absent). -struct emit_ctx -{ - // per-sample - const struct planner_sample* sample; - const struct zarr_metadata* meta; - const uint64_t* inner_per_shard_dim; // [meta->rank] - const uint64_t* chunk_beg; // [meta->rank], first chunk for sample - uint32_t sample_idx_in_batch; - uint16_t batch_pool_slot; - uint8_t codec_id; - uint32_t decompressed_n_bytes; // == inner_chunk_bytes(meta) (validated) - uint64_t page_alignment_bytes; - // Mutable handle to this sample's plan entry. emit_chunk lazily - // populates ->layout / ->layout_probed on the first non-fill emit. - struct sample_plan* sp; - const struct chunk_layout* chunk_layout; - uint32_t max_substreams_per_chunk; - // per-shard - const struct zarr_shard_entry* shard_entries; - uint64_t n_shard_entries; - const char* interned_path; - int shard_missing; -}; - -// Append a fill-mode chunk_plan + matching dummy read_op so the -// batch-level read_ops / chunk_plans arrays stay 1:1 with the -// chunk_plan stream that downstream input dispatch walks in lockstep. -static enum damacy_status -emit_fill_chunk(const struct emit_ctx* ctx, - const uint64_t* chunk_coord, - struct planner_output* out) -{ - if (out->n_read_ops >= out->read_ops_cap || - out->n_chunk_plans >= out->chunk_plans_cap) - return DAMACY_BUDGET; - const struct zarr_metadata* meta = ctx->meta; - - uint32_t read_op_idx = out->n_read_ops; - struct read_op* r = &out->read_ops[read_op_idx]; - *r = (struct read_op){ 0 }; - out->n_read_ops++; - - struct chunk_plan* cp = &out->chunk_plans[out->n_chunk_plans]; - *cp = (struct chunk_plan){ - .read_op_idx = read_op_idx, - .offset_in_read = 0, - .compressed_nbytes = 0, - .decompressed_nbytes = ctx->decompressed_n_bytes, - .batch_pool_slot = ctx->batch_pool_slot, - .sample_idx_in_batch = (uint16_t)ctx->sample_idx_in_batch, - .codec_id = (uint8_t)CODEC_FILL, - .is_fill = 1, - }; - for (uint8_t d = 0; d < meta->rank; ++d) - cp->chunk_d[d] = (uint32_t)(chunk_coord[d] - ctx->chunk_beg[d]); - out->n_chunk_plans++; - return DAMACY_OK; -} - -// Process one chunk: build read_op + chunk_plan, append to output. -// Empty shard entries and missing shard files emit a fill-mode chunk_plan -// referencing the array's fill_value rather than failing. -static enum damacy_status -emit_chunk(const struct emit_ctx* ctx, - const uint64_t* chunk_coord, - const uint64_t* local_inner, - struct planner_output* out) -{ - const struct zarr_metadata* meta = ctx->meta; - - // Shard file absent: every chunk in the shard is fill. - if (ctx->shard_missing) - return emit_fill_chunk(ctx, chunk_coord, out); - - uint64_t entry_idx = - row_major_linear(local_inner, ctx->inner_per_shard_dim, meta->rank); - if (entry_idx >= ctx->n_shard_entries) - return DAMACY_DECODE; - - const struct zarr_shard_entry* entry = &ctx->shard_entries[entry_idx]; - int off_sentinel = entry->offset == ZARR_SHARD_EMPTY_OFFSET; - int nb_sentinel = entry->nbytes == ZARR_SHARD_EMPTY_NBYTES; - if (off_sentinel && nb_sentinel) - return emit_fill_chunk(ctx, chunk_coord, out); - if (off_sentinel != nb_sentinel) { - log_error("zarr shard index entry %llu half-sentinel " - "(offset=0x%llx, nbytes=0x%llx); treating as corrupt", - (unsigned long long)entry_idx, - (unsigned long long)entry->offset, - (unsigned long long)entry->nbytes); - return DAMACY_DECODE; + if (self) { + dispatch_scratch_destroy(&self->scratch); + free(self); } - - if (entry->nbytes > DAMACY_MAX_CHUNK_BYTES) - return DAMACY_DECODE; - - // Layout was pre-fetched by the prefetcher into chunk_layout_cache. - // NULL means the array isn't blosc1 (decoder uses caps); the wave- - // eligibility gate keeps layout_probed=0 chunks out until probed. - if (ctx->sp && !ctx->sp->layout_probed && ctx->chunk_layout) { - ctx->sp->layout = *ctx->chunk_layout; - ctx->sp->layout_probed = 1; - } - - // Page-aligned read window enclosing [offset, offset + nbytes). - uint64_t page_alignment_bytes = ctx->page_alignment_bytes; - uint64_t aligned_file_offset = - align_down_u64(entry->offset, page_alignment_bytes); - uint64_t aligned_end = - align_up_u64(entry->offset + entry->nbytes, page_alignment_bytes); - uint64_t read_n_bytes = aligned_end - aligned_file_offset; - if (read_n_bytes > DAMACY_MAX_CHUNK_BYTES) - return DAMACY_DECODE; - uint32_t chunk_offset_in_read = - (uint32_t)(entry->offset - aligned_file_offset); - - if (out->n_read_ops >= out->read_ops_cap || - out->n_chunk_plans >= out->chunk_plans_cap) - return DAMACY_BUDGET; - - uint32_t read_op_idx = out->n_read_ops; - struct read_op* r = &out->read_ops[read_op_idx]; - r->shard_path = ctx->interned_path; - r->file_offset = aligned_file_offset; - r->nbytes = (uint32_t)read_n_bytes; - out->n_read_ops++; - - struct chunk_plan* cp = &out->chunk_plans[out->n_chunk_plans]; - *cp = (struct chunk_plan){ - .read_op_idx = read_op_idx, - .offset_in_read = chunk_offset_in_read, - .compressed_nbytes = (uint32_t)entry->nbytes, - .decompressed_nbytes = ctx->decompressed_n_bytes, - .batch_pool_slot = ctx->batch_pool_slot, - .sample_idx_in_batch = (uint16_t)ctx->sample_idx_in_batch, - .codec_id = ctx->codec_id, - }; - for (uint8_t d = 0; d < meta->rank; ++d) - cp->chunk_d[d] = (uint32_t)(chunk_coord[d] - ctx->chunk_beg[d]); - - out->n_chunk_plans++; - return DAMACY_OK; } enum damacy_status planner_plan_segment(struct planner* self, const struct planner_sample* samples, const struct planner_placement* placement, - const int64_t* dst_strides, - uint8_t dst_full_rank, - struct planner_output* out) + const int64_t* strides, + uint8_t rank, + struct dispatch_output* out) { - enum damacy_status status = DAMACY_OK; - - CHECK_SILENT(Invalid, self); - CHECK_SILENT(Invalid, samples); - CHECK_SILENT(Invalid, placement); - CHECK_SILENT(Invalid, out); - CHECK_SILENT(Invalid, out->read_ops); - CHECK_SILENT(Invalid, out->chunk_plans); - CHECK_SILENT(Invalid, out->sample_plans); - CHECK_SILENT(Invalid, out->paths); - CHECK_SILENT(Invalid, dst_strides); - CHECK_SILENT(Invalid, dst_full_rank >= 1); - path_intern_reset(out->paths); - out->n_read_ops = 0; - out->n_chunk_plans = 0; - out->n_sample_plans = 0; - - for (uint32_t sample_idx = 0; sample_idx < placement->n_samples; - ++sample_idx) { - uint32_t sample_idx_in_batch = - placement->sample_idx_begin_in_batch + sample_idx; - const struct planner_sample* sample = &samples[sample_idx]; - if (!sample->uri) { - status = DAMACY_INVAL; - goto Cleanup; - } - - const void* meta_value = NULL; - int meta_err = 0; - enum prefetch_state meta_state = prefetch_cache_query( - self->cfg.array_meta_cache, sample->h_meta, &meta_value, &meta_err); - if (meta_state == PREFETCH_STATE_PENDING) { - // Batch gate should make this unreachable. - log_error("planner: meta still PENDING (uri=%s)", sample->uri); - status = DAMACY_INVAL; - goto Cleanup; - } - if (meta_state == PREFETCH_STATE_ERROR) { - status = meta_err ? (enum damacy_status)meta_err : DAMACY_INVAL; - goto Cleanup; - } - const struct zarr_metadata* meta = (const struct zarr_metadata*)meta_value; - if (!meta) { - status = DAMACY_INVAL; - goto Cleanup; - } - if (sample->aabb.rank != meta->rank) { - status = DAMACY_RANK; - goto Cleanup; - } - if ((uint8_t)(meta->rank + 1) != dst_full_rank) { - status = DAMACY_RANK; - goto Cleanup; - } - if (!cast_path_supported(self->cfg.dst_dtype, meta->dtype)) { - status = DAMACY_DTYPE; - goto Cleanup; - } - uint64_t inner_per_shard_dim[DAMACY_MAX_RANK]; - if (zarr_metadata_inner_per_shard(meta, inner_per_shard_dim, NULL)) { - status = DAMACY_DECODE; - goto Cleanup; - } - - uint64_t decompressed_n_bytes = inner_chunk_bytes(meta); - if (decompressed_n_bytes == 0) { + if (!self || !samples || !placement || !strides || !out || !out->paths || + !out->read_ops || !out->chunk_plans || !out->sample_plans || + !placement->n_samples || placement->n_samples > UINT16_MAX || rank < 2 || + rank > DAMACY_MAX_RANK + 1) + return DAMACY_INVAL; + uint32_t count = placement->n_samples; + uint32_t begin = placement->sample_idx_begin_in_batch; + if (count > out->sample_plans_cap || begin > out->sample_plans_cap - count || + begin > UINT16_MAX - count) + return DAMACY_BUDGET; + struct damacy_batch_spec output = { .dtype = self->config.dst_dtype, + .sample_rank = rank - 1, + .samples_per_batch = count }; + for (uint8_t d = 0; d < output.sample_rank; ++d) { + int64_t lo = samples[0].aabb.dims[d].beg; + int64_t hi = samples[0].aabb.dims[d].end; + if (lo < 0 || hi <= lo) + return DAMACY_INVAL; + output.sample_shape[d] = hi - lo; + } + struct damacy_plan_limits limits = { + .max_chunks = out->chunk_plans_cap < out->read_ops_cap + ? out->chunk_plans_cap + : out->read_ops_cap, + .max_chunk_bytes = + self->config.max_chunk_uncompressed_bytes && + self->config.max_chunk_uncompressed_bytes < UINT32_MAX + ? (uint32_t)self->config.max_chunk_uncompressed_bytes + : UINT32_MAX, + .max_shards_per_sample = UINT32_MAX, + .max_plan_bytes = 64ull << 20 + }; + struct prepared_plan* plan = NULL; + enum damacy_status status = + prepared_plan_build(self->config.array_meta_cache, + self->config.shard_index_cache, + samples, + count, + &output, + &limits, + &plan); + if (status != DAMACY_OK) + return status; + for (uint32_t i = 0; i < plan->n_chunks; ++i) { + const struct plan_chunk* chunk = &plan->chunks[i]; + enum compression_codec codec = + plan->arrays[chunk->array].metadata.inner_codec.id; + if (!chunk->missing && codec != CODEC_NONE && codec != CODEC_ZSTD && + codec != CODEC_BLOSC_ZSTD) { status = DAMACY_DECODE; - goto Cleanup; - } - if (self->cfg.max_chunk_uncompressed_bytes > 0 && - decompressed_n_bytes > self->cfg.max_chunk_uncompressed_bytes) { - log_error("planner: chunk uncompressed=%llu exceeds runtime cap=%llu " - "(uri=%s)", - (unsigned long long)decompressed_n_bytes, - (unsigned long long)self->cfg.max_chunk_uncompressed_bytes, - sample->uri); - status = DAMACY_BUDGET; - goto Cleanup; - } - - uint64_t chunk_beg[DAMACY_MAX_RANK]; - uint64_t chunk_end[DAMACY_MAX_RANK]; - if (chunk_range(&sample->aabb, meta, chunk_beg, chunk_end)) { - status = DAMACY_INVAL; - goto Cleanup; - } - - if (sample_idx_in_batch >= out->sample_plans_cap) { - status = DAMACY_BUDGET; - goto Cleanup; - } - struct sample_plan* sp = &out->sample_plans[sample_idx_in_batch]; - *sp = (struct sample_plan){ - .batch_pool_slot = placement->batch_pool_slot, - .sample_idx_in_batch = (uint16_t)sample_idx_in_batch, - .rank = meta->rank, - .src_dtype = (uint8_t)meta->dtype, - .sample_dst_off_elems = (int64_t)sample_idx_in_batch * dst_strides[0], - .chunk_count = 0, - }; - memcpy(sp->fill_value, meta->fill_value, sizeof sp->fill_value); - int64_t src_strides[DAMACY_MAX_RANK]; - row_major_strides(meta->inner_chunk_shape, meta->rank, src_strides); - uint32_t chunk_count = 1; - for (uint8_t d = 0; d < meta->rank; ++d) { - uint32_t S = (uint32_t)meta->inner_chunk_shape[d]; - uint32_t N = (uint32_t)(chunk_end[d] - chunk_beg[d]); - int64_t chunk_grid_origin = (int64_t)(chunk_beg[d] * (uint64_t)S); - sp->dims[d] = (struct sample_dim){ - .chunk_shape = S, - .chunk_grid_extent = N, - .aabb_lo_relative = sample->aabb.dims[d].beg - chunk_grid_origin, - .aabb_extent = sample->aabb.dims[d].end - sample->aabb.dims[d].beg, - .dst_stride = dst_strides[d + 1], - .src_stride = src_strides[d], - }; - chunk_count *= N; - } - sp->chunk_count = chunk_count; - if (out->n_sample_plans < sample_idx_in_batch + 1u) - out->n_sample_plans = sample_idx_in_batch + 1u; - - const void* layout_value = NULL; - int layout_err = 0; - enum prefetch_state layout_state = - prefetch_cache_query(self->cfg.chunk_layout_cache, - sample->h_layout, - &layout_value, - &layout_err); - if (layout_state == PREFETCH_STATE_PENDING) { - // Batch gate should make this unreachable. - log_error("planner: chunk_layout still PENDING (uri=%s)", sample->uri); - status = DAMACY_INVAL; - goto Cleanup; - } - if (layout_state == PREFETCH_STATE_ERROR) { - status = layout_err ? (enum damacy_status)layout_err : DAMACY_DECODE; - goto Cleanup; - } - // NULL value on READY is legitimate for non-blosc codecs; - // emit_chunk falls back to worst-case caps. - const struct chunk_layout* layout = - (const struct chunk_layout*)layout_value; - - struct emit_ctx ctx = { - .sample = sample, - .meta = meta, - .inner_per_shard_dim = inner_per_shard_dim, - .chunk_beg = chunk_beg, - .sample_idx_in_batch = sample_idx_in_batch, - .batch_pool_slot = placement->batch_pool_slot, - .codec_id = (uint8_t)meta->inner_codec.id, - .decompressed_n_bytes = (uint32_t)decompressed_n_bytes, - .page_alignment_bytes = self->cfg.page_alignment, - .sp = sp, - .chunk_layout = layout, - .max_substreams_per_chunk = self->cfg.max_substreams_per_chunk, - }; - - struct sample_shard_iterator shard_it; - if (sample_shard_iterator_init(&shard_it, meta, &sample->aabb)) { - status = DAMACY_INVAL; - goto Cleanup; - } - - uint64_t shard_coord[DAMACY_MAX_RANK]; - uint32_t shard_idx_in_sample = 0; - while (sample_shard_iterator_next(&shard_it, shard_coord)) { - if (shard_idx_in_sample >= sample->n_shards) { - log_error( - "planner: shard count mismatch (uri=%s iterated=%u expected=%u)", - sample->uri, - shard_idx_in_sample + 1u, - sample->n_shards); - status = DAMACY_INVAL; - goto Cleanup; - } - struct prefetch_handle h = sample->h_shards[shard_idx_in_sample++]; - const struct shard_index_value* sv = - (const struct shard_index_value*)prefetch_cache_try_get( - self->cfg.shard_index_cache, h); - if (!sv) { - int err = 0; - enum prefetch_state st = - prefetch_cache_query(self->cfg.shard_index_cache, h, NULL, &err); - if (st == PREFETCH_STATE_PENDING) { - // Batch gate should have made this unreachable. - log_error("planner: shard still PENDING (uri=%s)", sample->uri); - status = DAMACY_INVAL; - goto Cleanup; - } - if (st == PREFETCH_STATE_READY) { - log_error("planner: shard READY with NULL value (uri=%s)", - sample->uri); - status = DAMACY_INVAL; - goto Cleanup; - } - if (err == DAMACY_NOTFOUND) { - ctx.shard_missing = 1; - ctx.shard_entries = NULL; - ctx.n_shard_entries = 0; - ctx.interned_path = NULL; - } else { - status = err ? (enum damacy_status)err : DAMACY_DECODE; - goto Cleanup; - } - } else { - ctx.shard_missing = 0; - ctx.shard_entries = sv->entries; - ctx.n_shard_entries = sv->n_entries; - if (zarr_shard_path_build( - &self->path_sb, sample->uri, shard_coord, meta->rank)) { - status = DAMACY_OOM; - goto Cleanup; - } - ctx.interned_path = - path_intern_acquire(out->paths, strbuf_cstr(&self->path_sb)); - if (!ctx.interned_path) { - status = DAMACY_OOM; - goto Cleanup; - } - } - - uint64_t chunk_beg_in_shard[DAMACY_MAX_RANK]; - uint64_t chunk_end_in_shard[DAMACY_MAX_RANK]; - for (uint8_t d = 0; d < meta->rank; ++d) { - uint64_t s_beg = shard_coord[d] * inner_per_shard_dim[d]; - uint64_t s_end = s_beg + inner_per_shard_dim[d]; - chunk_beg_in_shard[d] = s_beg > chunk_beg[d] ? s_beg : chunk_beg[d]; - chunk_end_in_shard[d] = s_end < chunk_end[d] ? s_end : chunk_end[d]; - } - - uint64_t chunk_coord[DAMACY_MAX_RANK]; - for (uint8_t d = 0; d < meta->rank; ++d) - chunk_coord[d] = chunk_beg_in_shard[d]; - - for (;;) { - uint64_t local_inner[DAMACY_MAX_RANK]; - for (uint8_t d = 0; d < meta->rank; ++d) - local_inner[d] = chunk_coord[d] % inner_per_shard_dim[d]; - - enum damacy_status emit_status = - emit_chunk(&ctx, chunk_coord, local_inner, out); - if (emit_status != DAMACY_OK) { - status = emit_status; - goto Cleanup; - } - - int finished = 1; - for (int d = (int)meta->rank - 1; d >= 0; --d) { - chunk_coord[d]++; - if (chunk_coord[d] < chunk_end_in_shard[d]) { - finished = 0; - break; - } - chunk_coord[d] = chunk_beg_in_shard[d]; - } - if (finished) - break; - } - } - if (shard_idx_in_sample != sample->n_shards) { - log_error( - "planner: shard count mismatch (uri=%s iterated=%u expected=%u)", - sample->uri, - shard_idx_in_sample, - sample->n_shards); - status = DAMACY_INVAL; - goto Cleanup; + goto Done; } } - - { - status = planner_ensure_scratch(self, out->n_read_ops); - if (status != DAMACY_OK) - goto Cleanup; - status = coalesce_chunks(out, - self->cfg.read_op_max_bytes, - self->cfg.max_chunks_per_wave, - self->scratch_u32, - self->scratch_ops); - if (status != DAMACY_OK) - goto Cleanup; - if (out->n_read_ops + 1u > 3u * self->scratch_u32_cap) { - status = DAMACY_BUDGET; - goto Cleanup; + status = dispatch_plan_build(plan, + placement->batch_pool_slot, + self->config.page_alignment, + self->config.read_op_max_bytes, + self->config.max_chunks_per_wave, + out, + &self->scratch); + if (status != DAMACY_OK) + goto Done; + memmove(out->sample_plans + begin, + out->sample_plans, + count * sizeof(*out->sample_plans)); + out->n_sample_plans = begin + count; + for (uint32_t i = 0; i < count; ++i) { + struct sample_plan* sample = &out->sample_plans[begin + i]; + sample->sample_idx_in_batch = (uint16_t)(begin + i); + sample->sample_dst_off_elems = (int64_t)(begin + i) * strides[0]; + for (uint8_t d = 0; d < output.sample_rank; ++d) + sample->dims[d].dst_stride = strides[d + 1]; + const void* layout = NULL; + int error = 0; + enum prefetch_state state = prefetch_cache_query( + self->config.chunk_layout_cache, samples[i].h_layout, &layout, &error); + if (state != PREFETCH_STATE_READY) { + status = error ? (enum damacy_status)error : DAMACY_INVAL; + goto Done; + } + if (layout) { + sample->layout = *(const struct chunk_layout*)layout; + sample->layout_probed = 1; } - status = - group_chunks_by_read(out, self->scratch_u32, self->scratch_chunk_plans); - if (status != DAMACY_OK) - goto Cleanup; } - - return DAMACY_OK; - -Cleanup: + for (uint32_t i = 0; i < out->n_chunk_plans; ++i) + out->chunk_plans[i].sample_idx_in_batch += (uint16_t)begin; +Done: + prepared_plan_destroy(plan); return status; - -Invalid: - return DAMACY_INVAL; } enum damacy_status planner_plan(struct planner* self, const struct planner_sample* samples, uint32_t n_samples, - uint16_t batch_pool_slot, - const int64_t* dst_strides, - uint8_t dst_full_rank, - struct planner_output* out) + uint16_t slot, + const int64_t* strides, + uint8_t rank, + struct dispatch_output* out) { - struct planner_placement placement = { - .batch_pool_slot = batch_pool_slot, - .sample_idx_begin_in_batch = 0, - .n_samples = n_samples, - }; return planner_plan_segment( - self, samples, &placement, dst_strides, dst_full_rank, out); + self, + samples, + &(struct planner_placement){ .batch_pool_slot = slot, + .n_samples = n_samples }, + strides, + rank, + out); } diff --git a/src/planner/planner.h b/src/planner/planner.h index 67ca8fe4..3006657d 100644 --- a/src/planner/planner.h +++ b/src/planner/planner.h @@ -1,20 +1,7 @@ -// Planner: samples × cached zarr metadata → per-chunk read jobs and -// transform records ready for the wave scheduler / IO pool / decompress -// / assemble pipeline. -// -// chunk_plan wave-scheduler fields (host_buf_offset, -// dev_decompressed_offset) are filled in by the scheduler; the planner -// zeroes them. #pragma once -#include "damacy.h" // damacy_status, damacy_sample, damacy_aabb -#include "damacy_limits.h" // DAMACY_MAX_RANK -#include "prefetch/prefetch_handle.h" -#include "zarr/zarr_chunk_layout.h" // struct chunk_layout -#include "zarr/zarr_metadata.h" // DAMACY_MAX_DTYPE_BYTES - -#include -#include +#include "executor/dispatch.h" +#include "planner/plan_builder.h" #ifdef __cplusplus extern "C" @@ -22,117 +9,6 @@ extern "C" #endif struct prefetch_cache; - struct path_intern; - - // Page-aligned IO operation. Multiple chunk_plans may share one - // read_op after coalescing. - // shard_path is interned by the planner; equal paths share a pointer, - // and the pointer is valid for the planner's lifetime — long enough - // for the wave scheduler's plan queue to outlive any batch. Fills set - // shard_path = NULL. - struct read_op - { - const char* shard_path; - uint64_t file_offset; // multiple of page_alignment - uint64_t host_buf_offset; // wave-scheduler-assigned; planner sets 0 - uint32_t nbytes; // multiple of page_alignment - }; - - // Per-dimension bundle for one sample. Co-locating all of dimension d's - // parameters lets the assemble kernel's R-loop touch one cache line per - // iteration instead of seven scattered loads. - struct sample_dim - { - uint32_t chunk_shape; // S[d] — uniform per source, no clipping - uint32_t chunk_grid_extent; // N[d] — chunks per dimension within sample - int64_t aabb_lo_relative; // aabb_lo[d] − chunk_grid_origin[d], in [0,S) - int64_t aabb_extent; // sample AABB extent - int64_t dst_stride; // batch tensor stride (elements) - int64_t src_stride; // chunk-local row-major stride (elements) - }; - - // Per-sample header consumed by assemble. All chunks within the sample - // share these constants (uniform shape, single source). Per-chunk - // records are reduced to {dev_decompressed_offset, chunk_d}. - struct sample_plan - { - uint16_t batch_pool_slot; - uint16_t sample_idx_in_batch; - uint8_t rank; // spatial rank - uint8_t src_dtype; // enum dtype; source zarr type - struct sample_dim dims[DAMACY_MAX_RANK]; // dims[0..rank) - int64_t sample_dst_off_elems; // sample slot start (elements; * dst bpe at - // runtime) - uint32_t chunk_count; // ∏ dims[d].chunk_grid_extent (informational; - // chunk_plans are not contiguous per-sample - // after group_chunks_by_read) - // Array-level fill_value (zarr v3 metadata). Chunks tagged is_fill - // broadcast these bytes; bytes are interpreted under src_dtype. - uint8_t fill_value[DAMACY_MAX_DTYPE_BYTES]; - // Per-array blosc1 chunk layout, populated lazily from the meta - // cache on first non-fill emit. layout_probed = 0 means the - // wave-eligibility gate rejects the wave before prepare_decode_caps - // is called. - struct chunk_layout layout; - uint8_t layout_probed; - }; - - // Per-chunk plan. Carries IO/decompress fields plus assemble-side - // chunk_d (grid position within sample, 0..N[d]) and sample_idx so - // the kernel can look up the sample_plan. - // - // is_fill marks chunks that are absent from the store (sparse zarr v3): - // read_op_idx / offset_in_read / compressed_nbytes are unused, the - // codec stage is skipped, and assemble broadcasts the sample's - // fill_value across the chunk's region instead of reading the arena. - // decompressed_nbytes is still set so wave accounting tracks the - // conceptual chunk size. - struct chunk_plan - { - uint32_t read_op_idx; - uint32_t offset_in_read; // chunk start within the read - uint32_t compressed_nbytes; - uint32_t decompressed_nbytes; - uint64_t host_buf_offset; // scheduler-assigned (per wave) - uint32_t dev_decompressed_offset; // scheduler-assigned (per wave) - uint16_t batch_pool_slot; - uint16_t sample_idx_in_batch; // index into planner_output.sample_plans - uint8_t codec_id; - uint8_t is_fill; // 1 = absent chunk; fill_value lives on the sample - uint32_t chunk_d[DAMACY_MAX_RANK]; // grid position within sample (0..N) - }; - - struct read_op_group - { - uint32_t read_op_idx; - uint32_t first_chunk; - uint32_t n_chunks; - uint64_t total_decompressed; - }; - - struct read_op_group_iterator - { - const struct read_op_group* groups; - uint32_t n_groups; - uint32_t cursor; - }; - - void read_op_group_iterator_init(struct read_op_group_iterator* it, - const struct read_op_group* groups, - uint32_t n_groups, - uint32_t start_group); - int read_op_group_iterator_next(struct read_op_group_iterator* it, - struct read_op_group* out); - - struct planner_sample - { - const char* uri; - struct damacy_aabb aabb; - struct prefetch_handle h_meta; - struct prefetch_handle* h_shards; - uint32_t n_shards; - struct prefetch_handle h_layout; - }; struct planner_placement { @@ -175,57 +51,13 @@ extern "C" struct planner** out); void planner_destroy(struct planner* p); - // Output buffers for planner_plan. Caller owns the storage; planner - // populates *_n on success. If any buffer fills before the plan - // completes, planner_plan returns DAMACY_BUDGET. - // - // `paths` interns each emitted read_op's shard_path. planner_plan - // resets it at entry; caller need not. - struct planner_output - { - struct read_op* read_ops; - uint32_t read_ops_cap; - uint32_t n_read_ops; - struct chunk_plan* chunk_plans; - uint32_t chunk_plans_cap; - uint32_t n_chunk_plans; - struct sample_plan* sample_plans; - uint32_t sample_plans_cap; - uint32_t n_sample_plans; - struct read_op_group* read_op_groups; - uint32_t read_op_groups_cap; - uint32_t n_read_op_groups; - struct path_intern* paths; - uint32_t n_chunks_to_load; // non-fill chunks (= IO requests pre-coalesce) - uint32_t n_loads_issued; // real (non-fill) read_ops after coalesce - }; - - // Plan one training batch — i.e., the N samples that land in one - // output tensor of shape [N, ...zarr_axes]. samples are processed in - // order; sample i becomes sample_plans[i].sample_idx_in_batch == i. - // All chunk_plans are tagged with batch_pool_slot so the scheduler - // knows which slot in the batch pool receives the assembled output. - // - // This is independent of the scheduler's wave granularity: the wave - // scheduler chunks the produced plan queue into wave-sized dispatch - // units (one batch typically spans multiple waves). The planner is - // batch-shaped; waves are a downstream concern. - // - // Empty inner chunks (offset == nbytes == 0xFFFF…) and shards that - // don't exist in the store are emitted as fill-mode chunk_plans - // carrying the array's fill_value; downstream skips IO/decompress and - // assemble broadcasts the fill bytes over the chunk's region. - // - // shard_path strings are interned by the planner; equal paths share - // a pointer across all emitted read_ops, and the storage lives until - // planner_destroy. enum damacy_status planner_plan(struct planner* p, const struct planner_sample* samples, uint32_t n_samples, uint16_t batch_pool_slot, const int64_t* dst_strides, // [rank+1] uint8_t dst_full_rank, // rank+1 - struct planner_output* out); + struct dispatch_output* out); enum damacy_status planner_plan_segment( struct planner* p, @@ -233,7 +65,7 @@ extern "C" const struct planner_placement* placement, const int64_t* dst_strides, // [rank+1] uint8_t dst_full_rank, // rank+1 - struct planner_output* out); + struct dispatch_output* out); #ifdef __cplusplus } diff --git a/src/prefetch/prefetcher.c b/src/prefetch/prefetcher.c index 16746bb6..b2a25448 100644 --- a/src/prefetch/prefetcher.c +++ b/src/prefetch/prefetcher.c @@ -56,7 +56,8 @@ struct prefetcher_slot struct prefetch_handle* h_shards; uint64_t* shard_coords; // flat [n_shards][rank] uint32_t n_shards; - uint32_t n_shards_requested; // resume cursor when a shard_index request AGAINs + uint32_t + n_shards_requested; // resume cursor when a shard_index request AGAINs struct prefetch_handle h_layout; }; @@ -365,6 +366,10 @@ advance_from_meta(struct prefetcher* p, struct prefetcher_slot* s) if (n == 0) { s->n_shards = 0; + if (!p->chunk_layout_cache) { + slot_mark_ready(p, s); + return; + } struct prefetch_request_result layout_req = request_chunk_layout(p, s); if (layout_req.status == DAMACY_AGAIN) return; @@ -430,6 +435,10 @@ advance_from_shard(struct prefetcher* p, struct prefetcher_slot* s) } } + if (!p->chunk_layout_cache) { + slot_mark_ready(p, s); + return; + } struct prefetch_request_result layout_req = request_chunk_layout(p, s); // AGAIN = the chunk_layout cache is transiently saturated. Stay in // pending_shards and retry next tick; the watermark advance will free a pin. @@ -624,7 +633,6 @@ prefetcher_create(const struct prefetcher_config* cfg) CHECK(Error, cfg->lookahead); CHECK(Error, cfg->array_meta_cache); CHECK(Error, cfg->shard_index_cache); - CHECK(Error, cfg->chunk_layout_cache); self = (struct prefetcher*)malloc(sizeof(*self)); CHECK(Error, self); @@ -934,7 +942,8 @@ prefetcher_advance_watermark(struct prefetcher* self, uint64_t watermark) return; prefetch_cache_advance_watermark(self->array_meta_cache, watermark); prefetch_cache_advance_watermark(self->shard_index_cache, watermark); - prefetch_cache_advance_watermark(self->chunk_layout_cache, watermark); + if (self->chunk_layout_cache) + prefetch_cache_advance_watermark(self->chunk_layout_cache, watermark); } void diff --git a/src/render_job/render_job.c b/src/render_job/render_job.c index 9e7d1df0..ac8d57b1 100644 --- a/src/render_job/render_job.c +++ b/src/render_job/render_job.c @@ -22,9 +22,12 @@ render_job_init(struct render_job* job, uint32_t samples_per_batch_cap) job->read_op_groups = (struct read_op_group*)calloc( DAMACY_MAX_CHUNKS_PER_BATCH, sizeof(struct read_op_group)); CHECK(Error, job->read_op_groups); - job->sample_plans = (struct sample_plan*)calloc(samples_per_batch_cap, - sizeof(struct sample_plan)); - CHECK(Error, job->sample_plans); + size_t sample_bytes = + (size_t)samples_per_batch_cap * sizeof(struct sample_plan); + if (cuMemHostAlloc((void**)&job->sample_plans, sample_bytes, 0) != + CUDA_SUCCESS) + goto Error; + memset(job->sample_plans, 0, sample_bytes); CUdeviceptr dptr = 0; if (cuMemAlloc(&dptr, (size_t)samples_per_batch_cap * sizeof(struct sample_plan)) != @@ -45,7 +48,8 @@ render_job_destroy(struct render_job* job, int cuda_skip) free(job->read_ops); free(job->chunk_plans); free(job->read_op_groups); - free(job->sample_plans); + if (!cuda_skip && job->sample_plans) + cuMemFreeHost(job->sample_plans); path_intern_free(&job->paths); if (!cuda_skip && job->d_sample_plans) cuMemFree(CUDPTR(job->d_sample_plans)); @@ -109,10 +113,10 @@ render_job_reset(struct render_job* job) job->n_groups_dispatched = 0; } -struct planner_output -render_job_planner_output(struct render_job* job, uint32_t samples_per_batch) +struct dispatch_output +render_job_dispatch_output(struct render_job* job, uint32_t samples_per_batch) { - return (struct planner_output){ + return (struct dispatch_output){ .read_ops = job->read_ops, .read_ops_cap = DAMACY_MAX_CHUNKS_PER_BATCH, .chunk_plans = job->chunk_plans, @@ -126,14 +130,15 @@ render_job_planner_output(struct render_job* job, uint32_t samples_per_batch) } enum damacy_status -render_job_upload_sample_plans(struct render_job* job) +render_job_upload_sample_plans(struct render_job* job, void* stream) { if (job->n_sample_plans == 0) return DAMACY_OK; - return cuMemcpyHtoD(CUDPTR(job->d_sample_plans), - job->sample_plans, - (size_t)job->n_sample_plans * - sizeof(struct sample_plan)) == CUDA_SUCCESS + return cuMemcpyHtoDAsync(CUDPTR(job->d_sample_plans), + job->sample_plans, + (size_t)job->n_sample_plans * + sizeof(struct sample_plan), + (CUstream)stream) == CUDA_SUCCESS ? DAMACY_OK : DAMACY_CUDA; } @@ -142,7 +147,7 @@ void render_job_commit_plan(struct render_job* job, uint16_t batch_pool_slot, uint64_t batch_id, - const struct planner_output* out) + const struct dispatch_output* out) { job->batch_pool_slot = batch_pool_slot; job->batch_id = batch_id; diff --git a/src/render_job/render_job.h b/src/render_job/render_job.h index b222759e..417b0417 100644 --- a/src/render_job/render_job.h +++ b/src/render_job/render_job.h @@ -8,7 +8,7 @@ #include "damacy.h" #include "damacy_limits.h" -#include "planner/planner.h" +#include "executor/dispatch.h" #include "util/path_intern.h" #include @@ -111,17 +111,17 @@ render_job_pool_get_const(const struct render_job_pool* pool, void render_job_reset(struct render_job* job); -struct planner_output -render_job_planner_output(struct render_job* job, uint32_t samples_per_batch); +struct dispatch_output +render_job_dispatch_output(struct render_job* job, uint32_t samples_per_batch); enum damacy_status -render_job_upload_sample_plans(struct render_job* job); +render_job_upload_sample_plans(struct render_job* job, void* stream); void render_job_commit_plan(struct render_job* job, uint16_t batch_pool_slot, uint64_t batch_id, - const struct planner_output* out); + const struct dispatch_output* out); int render_job_has_work(const struct render_job* job); diff --git a/src/scheduler/scheduler.c b/src/scheduler/scheduler.c index 27590d36..c5909fda 100644 --- a/src/scheduler/scheduler.c +++ b/src/scheduler/scheduler.c @@ -12,6 +12,7 @@ struct scheduler struct platform_mutex* m; struct platform_cond* cv; scheduler_step_fn step; + struct scheduler_hooks hooks; void* arg; int64_t idle_ns; int shutdown; // protected by m @@ -24,10 +25,18 @@ worker_main(void* p) { struct scheduler* s = (struct scheduler*)p; numa_apply_thread_affinity(&s->affinity, "scheduler_worker"); + if (s->hooks.enter) { + platform_mutex_lock(s->m); + s->hooks.enter(s->arg); + platform_cond_broadcast(s->cv); + platform_mutex_unlock(s->m); + } for (;;) { platform_mutex_lock(s->m); if (s->shutdown) { platform_mutex_unlock(s->m); + if (s->hooks.leave) + s->hooks.leave(s->arg); return; } int ready = s->step(s->arg); @@ -42,7 +51,8 @@ struct scheduler* scheduler_create(scheduler_step_fn step, void* arg, int64_t idle_ns, - const struct numa_resolved* affinity) + const struct numa_resolved* affinity, + const struct scheduler_hooks* hooks) { if (!step || idle_ns <= 0) { log_error("scheduler: invalid arguments (step=%d idle_ns=%lld)", @@ -56,6 +66,8 @@ scheduler_create(scheduler_step_fn step, return NULL; } s->step = step; + if (hooks) + s->hooks = *hooks; s->arg = arg; s->idle_ns = idle_ns; // Copy unconditionally; numa_apply_thread_affinity no-ops when @@ -81,17 +93,25 @@ scheduler_create(scheduler_step_fn step, return s; } +void +scheduler_stop(struct scheduler* s) +{ + if (!s || !s->thread) + return; + platform_mutex_lock(s->m); + s->shutdown = 1; + platform_cond_broadcast(s->cv); + platform_mutex_unlock(s->m); + platform_thread_join(s->thread); + s->thread = NULL; +} + void scheduler_destroy(struct scheduler* s) { if (!s) return; - if (s->thread) { - platform_mutex_lock(s->m); - s->shutdown = 1; - platform_mutex_unlock(s->m); - platform_thread_join(s->thread); - } + scheduler_stop(s); platform_cond_free(s->cv); platform_mutex_free(s->m); free(s); diff --git a/src/scheduler/scheduler.h b/src/scheduler/scheduler.h index a22d515b..a272d711 100644 --- a/src/scheduler/scheduler.h +++ b/src/scheduler/scheduler.h @@ -17,15 +17,23 @@ extern "C" // to have the scheduler broadcast its cond. typedef int (*scheduler_step_fn)(void* arg); + struct scheduler_hooks + { + void (*enter)(void*); + void (*leave)(void*); + }; + // Spawn the worker. idle_ns must be > 0. `affinity` is the resolved // NUMA placement plan from numa_init; pass NULL (or a struct with // node<0) to skip affinity. Returns NULL on failure. struct scheduler* scheduler_create(scheduler_step_fn step, void* arg, int64_t idle_ns, - const struct numa_resolved* affinity); + const struct numa_resolved* affinity, + const struct scheduler_hooks* hooks); // Signal shutdown, join, free. NULL-safe. + void scheduler_stop(struct scheduler* s); void scheduler_destroy(struct scheduler* s); void scheduler_lock(struct scheduler* s); diff --git a/src/threadpool/threadpool.c b/src/threadpool/threadpool.c index 00c123be..0fc62bb8 100644 --- a/src/threadpool/threadpool.c +++ b/src/threadpool/threadpool.c @@ -7,6 +7,7 @@ #include #include #include +#include #define THREADPOOL_SPIN_ITERS 10000u #define THREADPOOL_CACHELINE 64u @@ -122,8 +123,9 @@ threadpool_new(int nthreads) struct threadpool* p = NULL; CHECK_SILENT(Fail, nthreads >= 0); - p = (struct threadpool*)calloc(1, sizeof(*p)); + p = aligned_alloc(alignof(struct threadpool), sizeof(*p)); CHECK_SILENT(Fail, p); + memset(p, 0, sizeof(*p)); p->nworkers = nthreads; atomic_store_explicit(&p->epoch, 0, memory_order_relaxed); diff --git a/src/wave/wave_budget.c b/src/wave/wave_budget.c index a567ed5a..7d775a47 100644 --- a/src/wave/wave_budget.c +++ b/src/wave/wave_budget.c @@ -6,9 +6,10 @@ #include "decoder/blosc1.h" #include "decoder/decoder_memcpy.h" #include "decoder/decoder_zstd.h" +#include "executor/dispatch.h" #include "gpu_budget/gpu_budget.h" #include "log/log.h" -#include "planner/planner.h" +#include "platform/platform.h" #include "util/cuda_check.h" #include "wave/fanout.h" // fanout_next_pow2 @@ -207,6 +208,7 @@ wave_pool_resolve_sizing(uint32_t max_chunks_per_wave, uint32_t max_substreams_per_chunk, uint8_t input_device_staging_buffers, uint64_t max_gpu_memory_bytes, + uint64_t max_read_op_bytes, uint64_t max_chunk_uncompressed_bytes, uint32_t samples_per_batch, struct wave_pool_sizing* out) @@ -214,11 +216,14 @@ wave_pool_resolve_sizing(uint32_t max_chunks_per_wave, const uint32_t max_substreams_per_wave = DAMACY_MAX_SUBSTREAMS_PER_WAVE( max_chunks_per_wave, max_substreams_per_chunk); const uint64_t min_per_wave = max_chunk_uncompressed_bytes; + const uint64_t alignment = platform_page_alignment(); + const uint64_t min_input = + min_per_wave > alignment ? min_per_wave : alignment; uint64_t total_min = 0; enum damacy_status s = predict_pool_total(max_chunks_per_wave, max_substreams_per_wave, input_device_staging_buffers, - min_per_wave, + min_input, min_per_wave, max_chunk_uncompressed_bytes, samples_per_batch, @@ -251,7 +256,7 @@ wave_pool_resolve_sizing(uint32_t max_chunks_per_wave, s = predict_pool_total(max_chunks_per_wave, max_substreams_per_wave, input_device_staging_buffers, - per_wave, + per_wave > min_input ? per_wave : min_input, per_wave, max_chunk_uncompressed_bytes, samples_per_batch, @@ -265,7 +270,7 @@ wave_pool_resolve_sizing(uint32_t max_chunks_per_wave, s = predict_pool_total(max_chunks_per_wave, max_substreams_per_wave, input_device_staging_buffers, - per_wave, + per_wave > min_input ? per_wave : min_input, per_wave, max_chunk_uncompressed_bytes, samples_per_batch, @@ -273,7 +278,16 @@ wave_pool_resolve_sizing(uint32_t max_chunks_per_wave, if (s != DAMACY_OK) return s; } - out->input_staging_per_wave = per_wave; + uint64_t input_bytes = per_wave > min_input ? per_wave : min_input; + if (input_device_staging_buffers && input_bytes < max_read_op_bytes) { + uint64_t extra = + (max_gpu_memory_bytes - predicted) / input_device_staging_buffers; + if (extra > max_read_op_bytes - input_bytes) + extra = max_read_op_bytes - input_bytes; + input_bytes += extra; + predicted += extra * input_device_staging_buffers; + } + out->input_staging_per_wave = input_bytes; out->dev_decompressed_per_wave = per_wave; out->worst_case_total_bytes = predicted; return DAMACY_OK; diff --git a/src/wave/wave_budget.h b/src/wave/wave_budget.h index 2f74e35b..b0b1adc1 100644 --- a/src/wave/wave_budget.h +++ b/src/wave/wave_budget.h @@ -64,6 +64,7 @@ wave_pool_resolve_sizing(uint32_t max_chunks_per_wave, uint32_t max_substreams_per_chunk, uint8_t input_device_staging_buffers, uint64_t max_gpu_memory_bytes, + uint64_t max_read_op_bytes, uint64_t max_chunk_uncompressed_bytes, uint32_t samples_per_batch, struct wave_pool_sizing* out); diff --git a/src/wave/wave_pool.c b/src/wave/wave_pool.c index 8e99bae3..5daf22ec 100644 --- a/src/wave/wave_pool.c +++ b/src/wave/wave_pool.c @@ -6,11 +6,11 @@ #include "decoder/blosc1_parse.h" #include "decoder/decoder_zstd.h" #include "decoder/status_reduce.h" +#include "executor/dispatch.h" #include "fanout.h" #include "gpu_budget/gpu_budget.h" #include "log/log.h" #include "nvtx/nvtx.h" -#include "planner/planner.h" #include "render_job/render_job.h" #include "store/store.h" #include "util/cuda_check.h" diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b8e5ed6d..bce68503 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -45,9 +45,9 @@ add_damacy_test( path_intern test_fixture ) -add_damacy_test(test_coalesce planner path_intern log) +add_damacy_test(test_coalesce dispatch_utils path_intern log) add_damacy_test(test_path_intern path_intern hash log) -add_damacy_test(test_group_chunks planner log) +add_damacy_test(test_group_chunks dispatch_utils log) add_damacy_test(test_lookahead lookahead log platform) add_damacy_test(test_prefetch_cache prefetch_cache log) add_damacy_test(test_array_meta array_meta prefetch_cache zarr store test_fixture) @@ -62,6 +62,12 @@ add_damacy_test(test_scheduler scheduler platform log) # DAMACY_FUZZ (which disables the CUDA-linked damacy target). Also # requires `uv` on PATH at test time (provided by flake.nix devShell). if(NOT DAMACY_FUZZ) + add_damacy_test(test_damacy_plan damacy) + add_damacy_test(test_cpu_pipeline damacy test_fixture Threads::Threads) + set_tests_properties(test_cpu_pipeline PROPERTIES TIMEOUT 120) +endif() + +if(DAMACY_CUDA) add_library(cuda_test_util STATIC cuda_init.c cuda_init.h) target_include_directories( cuda_test_util @@ -76,7 +82,6 @@ if(NOT DAMACY_FUZZ) test_damacy PRIVATE CUDA::cudart_static cuda_test_util ) - add_damacy_test(test_damacy_plan damacy) add_damacy_test(test_damacy_blosc damacy test_fixture) target_link_libraries(test_damacy_blosc PRIVATE cuda_test_util) add_damacy_test(test_damacy_caps damacy decoder test_fixture cuda_test_util) diff --git a/tests/test_assemble.c b/tests/test_assemble.c index c07cb919..01ff2827 100644 --- a/tests/test_assemble.c +++ b/tests/test_assemble.c @@ -6,8 +6,8 @@ #include "assemble/assemble.h" #include "dtype/dtype.h" +#include "executor/dispatch.h" #include "expect.h" -#include "planner/planner.h" #include #include diff --git a/tests/test_coalesce.c b/tests/test_coalesce.c index 8a30bf55..70af6d19 100644 --- a/tests/test_coalesce.c +++ b/tests/test_coalesce.c @@ -1,11 +1,11 @@ -// Unit tests for planner/coalesce.c: synthetic planner_output → sort + +// Unit tests for executor/coalesce.c: synthetic dispatch_output → sort + // fuse-with-cap + interleave. No zarr/store/cache plumbing — every // input read_op and chunk_plan is constructed inline. #include "damacy_limits.h" +#include "executor/coalesce.h" +#include "executor/dispatch.h" #include "expect.h" -#include "planner/coalesce.h" -#include "planner/planner.h" #include "util/path_intern.h" #include @@ -57,7 +57,7 @@ mk_fill(struct read_op* r, struct chunk_plan* cp, uint32_t read_op_idx) // Allocate scratch sized for n input read_ops and call coalesce. static enum damacy_status -run_coalesce(struct planner_output* out, uint64_t cap, uint32_t n_in) +run_coalesce(struct dispatch_output* out, uint64_t cap, uint32_t n_in) { uint32_t* u32 = (uint32_t*)calloc((size_t)n_in * 4u, sizeof(uint32_t)); struct read_op* ops = (struct read_op*)calloc(n_in, sizeof(struct read_op)); @@ -72,7 +72,7 @@ run_coalesce(struct planner_output* out, uint64_t cap, uint32_t n_in) static int test_empty(void) { - struct planner_output out = { 0 }; + struct dispatch_output out = { 0 }; struct read_op r = { 0 }; struct chunk_plan c = { 0 }; out.read_ops = &r; @@ -94,7 +94,7 @@ test_single_chunk_passthrough(void) struct read_op reads[1] = { 0 }; struct chunk_plan chunks[1] = { 0 }; mk(&reads[0], &chunks[0], "shard/0", 4096, 4096, 0, 100); - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 1, .n_read_ops = 1, @@ -121,7 +121,7 @@ test_touching_fuses(void) struct chunk_plan chunks[2] = { 0 }; mk(&reads[0], &chunks[0], "shard/0", 0, 4096, 0, 50); mk(&reads[1], &chunks[1], "shard/0", 4096, 4096, 1, 200); - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 2, .n_read_ops = 2, @@ -150,7 +150,7 @@ test_overlapping_fuses(void) struct chunk_plan chunks[2] = { 0 }; mk(&reads[0], &chunks[0], "shard/0", 0, 8192, 0, 100); mk(&reads[1], &chunks[1], "shard/0", 4096, 4096, 1, 50); - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 2, .n_read_ops = 2, @@ -175,7 +175,7 @@ test_gap_no_fusion(void) struct chunk_plan chunks[2] = { 0 }; mk(&reads[0], &chunks[0], "shard/0", 0, 4096, 0, 0); mk(&reads[1], &chunks[1], "shard/0", 8192, 4096, 1, 0); - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 2, .n_read_ops = 2, @@ -197,7 +197,7 @@ test_different_paths_no_fusion(void) struct chunk_plan chunks[2] = { 0 }; mk(&reads[0], &chunks[0], "shard/0", 0, 4096, 0, 0); mk(&reads[1], &chunks[1], "shard/1", 0, 4096, 1, 0); - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 2, .n_read_ops = 2, @@ -228,7 +228,7 @@ test_cap_splits(void) 4096, i, /*offset_in_read*/ 0); - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = N, .n_read_ops = N, @@ -257,7 +257,7 @@ test_single_over_cap(void) struct read_op reads[1] = { 0 }; struct chunk_plan chunks[1] = { 0 }; mk(&reads[0], &chunks[0], "shard/0", 0, 1u << 20, 0, 0); // 1 MB - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 1, .n_read_ops = 1, @@ -288,7 +288,7 @@ test_non_overlapping_output(void) mk(&reads[2], &chunks[2], "shard/A", 0, 4096, 2, 0); mk(&reads[3], &chunks[3], "shard/B", 0, 4096, 3, 0); mk(&reads[4], &chunks[4], "shard/A", 4096, 4096, 4, 0); - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = N, .n_read_ops = N, @@ -325,7 +325,7 @@ test_round_robin_interleave(void) mk(&reads[2], &chunks[2], "shard/A", 16384, 4096, 2, 0); mk(&reads[3], &chunks[3], "shard/B", 0, 4096, 3, 0); mk(&reads[4], &chunks[4], "shard/B", 8192, 4096, 4, 0); - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = N, .n_read_ops = N, @@ -345,8 +345,7 @@ test_round_robin_interleave(void) EXPECT(strcmp(reads[3].shard_path, "shard/B") == 0); EXPECT(reads[3].file_offset == 8192 && reads[3].nbytes == 4096); EXPECT(chunks[0].read_op_idx == 0 && chunks[0].offset_in_read == 50); - EXPECT(chunks[1].read_op_idx == 0 && - chunks[1].offset_in_read == 200 + 4096); + EXPECT(chunks[1].read_op_idx == 0 && chunks[1].offset_in_read == 200 + 4096); EXPECT(chunks[2].read_op_idx == 2); EXPECT(chunks[3].read_op_idx == 1); EXPECT(chunks[4].read_op_idx == 3); @@ -368,7 +367,7 @@ test_chunk_count_cap(void) EXPECT(reads && chunks); for (uint32_t i = 0; i < n; ++i) mk(&reads[i], &chunks[i], "shard/0", (uint64_t)i * step, step, i, 0); - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = n, .n_read_ops = n, @@ -402,7 +401,7 @@ test_fills_passthrough(void) mk(&reads[0], &chunks[0], "shard/0", 0, 4096, 0, 0); mk_fill(&reads[1], &chunks[1], 1); mk(&reads[2], &chunks[2], "shard/0", 4096, 4096, 2, 0); - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 3, .n_read_ops = 3, diff --git a/tests/test_cpu_pipeline.c b/tests/test_cpu_pipeline.c new file mode 100644 index 00000000..e5cd1c8a --- /dev/null +++ b/tests/test_cpu_pipeline.c @@ -0,0 +1,354 @@ +#include "damacy_pipeline.h" +#include "fixture.h" +#include "pipeline/components.h" +#include "platform/platform.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +struct components +{ + struct damacy_reader* reader; + struct damacy_metadata_reader* metadata_reader; + struct damacy_metadata* metadata; + struct damacy_planner* planner; + struct damacy_executor* executor; + struct damacy* pipeline; +}; + +static void +destroy_components(struct components* c) +{ + damacy_destroy(c->pipeline); + damacy_executor_destroy(c->executor); + damacy_planner_destroy(c->planner); + damacy_metadata_destroy(c->metadata); + damacy_metadata_reader_destroy(c->metadata_reader); + damacy_reader_destroy(c->reader); + *c = (struct components){ 0 }; +} + +static int +create_components(struct components* c) +{ + EXPECT(damacy_file_reader_create(2, 4, &c->reader) == DAMACY_OK); + EXPECT(damacy_file_metadata_reader_create(4, NULL, &c->metadata_reader) == + DAMACY_OK); + EXPECT( + damacy_zarr_metadata_create(c->metadata_reader, + &(struct damacy_metadata_cache_config){ + .array_entries = 8, .shard_entries = 32 }, + &c->metadata) == DAMACY_OK); + EXPECT(damacy_chunk_planner_create( + c->metadata, + &(struct damacy_plan_limits){ .max_chunks = 1024, + .max_chunk_bytes = 1 << 20, + .max_shards_per_sample = 4, + .max_plan_bytes = 1 << 20 }, + &c->planner) == DAMACY_OK); + EXPECT(damacy_cpu_executor_create(c->reader, + &(struct damacy_cpu_config){ + .decode_workers = 2, + .max_encoded_chunk_bytes = (1 << 20) + 1, + .max_decoded_chunk_bytes = 1 << 20, + .max_memory_bytes = 32 << 20 }, + &c->executor) == DAMACY_OK); + return 0; +} + +static int +start_pipeline(struct components* c, + enum damacy_dtype dtype, + int rows, + int cols, + uint32_t samples) +{ + EXPECT(damacy_pipeline_create( + c->planner, + c->executor, + &(struct damacy_batch_spec){ .dtype = dtype, + .sample_shape = { rows, cols }, + .sample_rank = 2, + .samples_per_batch = samples }, + &(struct damacy_queue_limits){ .lookahead_samples = 4, + .prepared_batches = 2 }, + &c->pipeline) == DAMACY_OK); + return 0; +} + +static struct damacy_sample +sample(const char* uri, int y, int x, int rows, int cols) +{ + return (struct damacy_sample){ + .uri = uri, + .aabb = { .rank = 2, .dims = { { y, y + rows }, { x, x + cols } } } + }; +} + +static int +verify_crop(struct damacy_batch* batch, int offset, int unsigned_bits) +{ + struct damacy_batch_info info; + damacy_batch_info(batch, &info); + EXPECT(info.device_type == DAMACY_DEVICE_CPU); + EXPECT(info.device_id == 0); + EXPECT(info.data == info.device_ptr); + EXPECT(info.ready_stream == NULL); + EXPECT(info.rank == 3 && info.shape[0] == 2 && info.shape[1] == 2 && + info.shape[2] == 5); + const float* values = info.data; + for (unsigned s = 0; s < 2; ++s) + for (int y = 0; y < 2; ++y) + for (int x = 0; x < 5; ++x) { + int64_t expected = (y + 1) * 11 + x + 2 + offset; + if (unsigned_bits) + expected &= ((1ll << unsigned_bits) - 1); + EXPECT(values[s * 10 + y * 5 + x] == (float)expected); + } + return 0; +} + +static int +test_codecs_and_types(void) +{ + const char* codecs[] = { "none", "zstd", "blosc-zstd" }; + const char* types[] = { "uint8", "uint16", "int16", "uint32", + "int32", "float16", "float32" }; + const int unsigned_bits[] = { 8, 16, 0, 32, 0, 0, 0 }; + for (unsigned ci = 0; ci < 3; ++ci) + for (unsigned ti = 0; ti < 7; ++ti) { + char root[] = "/tmp/damacy_cpu_XXXXXX"; + EXPECT(mkdtemp(root)); + char uri[256]; + snprintf(uri, sizeof(uri), "%s/array", root); + int64_t shape[] = { 5, 11 }, chunks[] = { 2, 4 }, shards[] = { 4, 8 }; + EXPECT(fixture_write_zarr_codec( + uri, shape, chunks, shards, 2, types[ti], -20, codecs[ci]) == 0); + struct components c = { 0 }; + EXPECT(create_components(&c) == 0); + EXPECT(start_pipeline(&c, DAMACY_F32, 2, 5, 2) == 0); + struct damacy_sample samples[] = { sample(uri, 1, 2, 2, 5), + sample(uri, 1, 2, 2, 5) }; + EXPECT(damacy_push(c.pipeline, + (struct damacy_sample_slice){ samples, samples + 2 }) + .status == DAMACY_OK); + struct damacy_batch* batch = NULL; + EXPECT(damacy_pop(c.pipeline, &batch) == DAMACY_OK); + EXPECT(verify_crop(batch, -20, unsigned_bits[ti]) == 0); + struct damacy_stats stats; + damacy_stats_get(c.pipeline, &stats); + EXPECT(stats.chunks_planned == 8); + EXPECT(stats.chunks_dispatched == 4); + EXPECT(stats.assemble.output_bytes == 2 * 2 * 5 * sizeof(float)); + EXPECT(stats.host_bytes_committed <= (32u << 20)); + EXPECT(stats.gpu_bytes_committed == 0); + damacy_release(c.pipeline, batch); + destroy_components(&c); + fixture_rm_tree(root); + } + return 0; +} + +static int +test_owned_plan(void) +{ + char root[] = "/tmp/damacy_owned_plan_XXXXXX"; + EXPECT(mkdtemp(root)); + char uri[256]; + snprintf(uri, sizeof(uri), "%s/array", root); + int64_t shape[] = { 5, 11 }, chunks[] = { 2, 4 }, shards[] = { 4, 8 }; + EXPECT(fixture_write_zarr(uri, shape, chunks, shards, 2, "uint16", 10) == 0); + struct components c = { 0 }; + EXPECT(create_components(&c) == 0); + struct damacy_batch_spec output = { .dtype = DAMACY_F32, + .sample_shape = { 2, 5 }, + .sample_rank = 2, + .samples_per_batch = 2 }; + struct damacy_queue_limits queues = { .lookahead_samples = 4, + .prepared_batches = 2 }; + EXPECT(c.planner->ops->start(c.planner, &output, &queues) == DAMACY_OK); + struct prepared_plan* saved = NULL; + for (unsigned i = 0; i < 12; ++i) { + struct damacy_sample samples[] = { sample(uri, 1, 2, 2, 5), + sample(uri, 1, 2, 2, 5) }; + EXPECT( + c.planner->ops + ->push(c.planner, (struct damacy_sample_slice){ samples, samples + 2 }) + .status == DAMACY_OK); + struct prepared_plan* plan = NULL; + enum damacy_status status = DAMACY_AGAIN; + for (unsigned retry = 0; retry < 10000 && status == DAMACY_AGAIN; ++retry) { + status = c.planner->ops->next(c.planner, &plan); + if (status == DAMACY_AGAIN) + platform_sleep_ns(1000000); + } + EXPECT(status == DAMACY_OK); + if (!saved) + saved = plan; + else + prepared_plan_destroy(plan); + strcat(uri, "/."); + } + EXPECT(saved->n_arrays == 1 && saved->n_chunks == 4 && saved->n_uses == 8); + damacy_planner_destroy(c.planner); + c.planner = NULL; + damacy_metadata_destroy(c.metadata); + c.metadata = NULL; + damacy_metadata_reader_destroy(c.metadata_reader); + c.metadata_reader = NULL; + EXPECT(saved->arrays[0].metadata.shape[1] == 11); + struct damacy_stats stats = { 0 }; + EXPECT(c.executor->ops->start(c.executor, &output, &stats) == DAMACY_OK); + EXPECT(c.executor->ops->submit(c.executor, saved, 42) == DAMACY_OK); + struct damacy_batch* batch = NULL; + for (unsigned retry = 0; retry < 10000 && !batch; ++retry) { + int changed = 0; + EXPECT(c.executor->ops->step(c.executor, &changed) == DAMACY_OK); + enum damacy_status status = c.executor->ops->take(c.executor, &batch); + EXPECT(status == DAMACY_OK || status == DAMACY_AGAIN); + if (!batch) + platform_sleep_ns(1000000); + } + EXPECT(batch); + EXPECT(verify_crop(batch, 10, 16) == 0); + damacy_batch_release(batch); + destroy_components(&c); + fixture_rm_tree(root); + return 0; +} + +struct pop_waiter +{ + struct damacy* pipeline; + _Atomic int started; + enum damacy_status status; +}; + +static void* +wait_pop(void* arg) +{ + struct pop_waiter* waiter = arg; + atomic_store(&waiter->started, 1); + struct damacy_batch* batch = NULL; + waiter->status = damacy_pop(waiter->pipeline, &batch); + damacy_batch_release(batch); + return NULL; +} + +static int +test_retained_outputs_and_shutdown(void) +{ + char root[] = "/tmp/damacy_lifetime_XXXXXX"; + EXPECT(mkdtemp(root)); + char uri[256]; + snprintf(uri, sizeof(uri), "%s/array", root); + int64_t shape[] = { 5, 11 }, chunks[] = { 2, 4 }, shards[] = { 4, 8 }; + EXPECT(fixture_write_zarr(uri, shape, chunks, shards, 2, "uint16", 10) == 0); + struct components c = { 0 }; + EXPECT(create_components(&c) == 0); + EXPECT(start_pipeline(&c, DAMACY_F32, 2, 5, 2) == 0); + struct damacy_sample samples[4]; + for (unsigned i = 0; i < 4; ++i) + samples[i] = sample(uri, 1, 2, 2, 5); + EXPECT(damacy_push(c.pipeline, + (struct damacy_sample_slice){ samples, samples + 4 }) + .status == DAMACY_OK); + struct damacy_batch *first = NULL, *second = NULL; + EXPECT(damacy_pop(c.pipeline, &first) == DAMACY_OK); + EXPECT(damacy_pop(c.pipeline, &second) == DAMACY_OK); + EXPECT(first != second); + struct damacy_batch_info a, b; + damacy_batch_info(first, &a); + damacy_batch_info(second, &b); + EXPECT(a.batch_id == 0 && b.batch_id == 1 && a.data != b.data); + damacy_batch_retain(first); + damacy_release(c.pipeline, first); + EXPECT(damacy_push(c.pipeline, + (struct damacy_sample_slice){ samples, samples + 2 }) + .status == DAMACY_OK); + struct pop_waiter waiter = { .pipeline = c.pipeline }; + pthread_t thread; + EXPECT(pthread_create(&thread, NULL, wait_pop, &waiter) == 0); + while (!atomic_load(&waiter.started)) + platform_sleep_ns(1000000); + damacy_shutdown(c.pipeline); + EXPECT(pthread_join(thread, NULL) == 0); + EXPECT(waiter.status == DAMACY_SHUTDOWN); + destroy_components(&c); + EXPECT(verify_crop(first, 10, 16) == 0); + EXPECT(verify_crop(second, 10, 16) == 0); + damacy_batch_release(first); + damacy_batch_release(second); + fixture_rm_tree(root); + return 0; +} + +static int +test_bfloat_rounding_and_fill(void) +{ + char root[] = "/tmp/damacy_bfloat_XXXXXX"; + EXPECT(mkdtemp(root)); + char path[256]; + snprintf(path, sizeof(path), "%s/zarr.json", root); + EXPECT(fixture_write_file( + path, + "{\"zarr_format\":3,\"node_type\":\"array\",\"shape\":[1,8]," + "\"data_type\":\"float32\",\"fill_value\":-2," + "\"chunk_grid\":{\"name\":\"regular\",\"configuration\":{\"chunk_" + "shape\":[1,8]}}," + "\"chunk_key_encoding\":{\"name\":\"default\",\"configuration\":{" + "\"separator\":\"/\"}}," + "\"codecs\":[{\"name\":\"bytes\",\"configuration\":{\"endian\":" + "\"little\"}}]}") == 0); + snprintf(path, sizeof(path), "%s/c", root); + EXPECT(mkdir(path, 0700) == 0); + snprintf(path, sizeof(path), "%s/c/0", root); + EXPECT(mkdir(path, 0700) == 0); + snprintf(path, sizeof(path), "%s/c/0/0", root); + uint32_t source[] = { 0x3f808000, 0x3f818000, 0xbf808000, 0x80000000, + 0x7f800000, 0xff800000, 0x7fc00000, 0x00008000 }; + uint16_t expected[] = { 0x3f80, 0x3f82, 0xbf80, 0x8000, + 0x7f80, 0xff80, 0x7fff, 0 }; + int fd = open(path, O_WRONLY | O_CREAT, 0600); + EXPECT(fd >= 0 && + write(fd, source, sizeof(source)) == (ssize_t)sizeof(source)); + close(fd); + for (unsigned missing = 0; missing < 2; ++missing) { + struct components c = { 0 }; + EXPECT(create_components(&c) == 0); + EXPECT(start_pipeline(&c, DAMACY_BF16, 1, 8, 1) == 0); + struct damacy_sample request = sample(root, 0, 0, 1, 8); + EXPECT(damacy_push(c.pipeline, + (struct damacy_sample_slice){ &request, &request + 1 }) + .status == DAMACY_OK); + struct damacy_batch* batch = NULL; + EXPECT(damacy_pop(c.pipeline, &batch) == DAMACY_OK); + struct damacy_batch_info info; + damacy_batch_info(batch, &info); + for (unsigned i = 0; i < 8; ++i) + EXPECT(((const uint16_t*)info.data)[i] == + (missing ? 0xc000 : expected[i])); + damacy_batch_release(batch); + destroy_components(&c); + if (!missing) + EXPECT(unlink(path) == 0); + } + fixture_rm_tree(root); + return 0; +} + +int +main(void) +{ + RUN(test_codecs_and_types); + RUN(test_owned_plan); + RUN(test_retained_outputs_and_shutdown); + RUN(test_bfloat_rounding_and_fill); + return 0; +} diff --git a/tests/test_damacy.c b/tests/test_damacy.c index 5014571a..e8b5580a 100644 --- a/tests/test_damacy.c +++ b/tests/test_damacy.c @@ -22,7 +22,6 @@ #include "cuda_init.h" #include "damacy.h" -#include "damacy_internal.h" #include "fixture.h" #include "platform/platform.h" #include "spin_kernel.h" @@ -113,7 +112,8 @@ run_one(struct damacy* d, damacy_batch_info(b, &info); EXPECT(info.rank == 3); EXPECT(info.dtype == DAMACY_F32); - EXPECT(info.ready_stream == (void*)d->wave_pool.stream_post); + EXPECT(info.ready_stream != NULL); + EXPECT(info.device_type == DAMACY_DEVICE_CUDA); EXPECT(info.shape[0] == 1); size_t n_elements = (size_t)info.shape[1] * (size_t)info.shape[2]; EXPECT(n_elements <= out_capacity_elements); diff --git a/tests/test_damacy_blosc.c b/tests/test_damacy_blosc.c index 8a39adcf..efc88971 100644 --- a/tests/test_damacy_blosc.c +++ b/tests/test_damacy_blosc.c @@ -284,13 +284,6 @@ test_multi_wave_per_batch(void) 0); } - // Tight max_gpu_memory_bytes + small chunk cap drives the resolver - // to pick the minimum per-wave geometry (one chunk per wave). With - // max_chunk_uncompressed_bytes = 4 KiB and a budget just barely big - // enough to fit total_min, dev_decompressed_per_wave lands near - // 4 KiB. input_dispatch_wave's per-chunk read_op is page-aligned (typically - // 4 KiB), so the 16-chunk batch spills into ≥2 waves of the same - // batch slot. struct damacy_config cfg = { .samples_per_batch = 4, .lookahead_samples = 8, @@ -307,8 +300,8 @@ test_multi_wave_per_batch(void) cfg.tuning.n_chunk_layout_cache = 16; cfg.tuning.max_shards_per_sample = 1; cfg.tuning.max_chunk_uncompressed_bytes = 4ull << 10; - // Resolver minimum so the 16-chunk batch spills into ≥2 waves. cfg.tuning.max_gpu_memory_bytes = 116ull << 20; + cfg.tuning.max_chunks_per_wave = 4; struct damacy* d = NULL; EXPECT(damacy_create(&cfg, &d) == DAMACY_OK); diff --git a/tests/test_damacy_plan.c b/tests/test_damacy_plan.c index 66501a6f..4debb4fb 100644 --- a/tests/test_damacy_plan.c +++ b/tests/test_damacy_plan.c @@ -2,47 +2,130 @@ #include "damacy_stats.h" #include "expect.h" -#include +#include + +struct test_planner +{ + struct damacy_planner base; + unsigned remaining; + unsigned prepared; + enum damacy_status failure; +}; + +struct test_executor +{ + struct damacy_executor base; + enum damacy_status status; + unsigned accepted; + unsigned submissions; + uint64_t ids[4]; + struct prepared_plan* last; +}; + +static enum damacy_status +prepare(struct damacy_planner* base, struct prepared_plan** out) +{ + struct test_planner* self = (void*)base; + if (self->failure) + return self->failure; + if (!self->remaining) + return DAMACY_AGAIN; + *out = calloc(1, sizeof(**out)); + if (!*out) + return DAMACY_OOM; + --self->remaining; + ++self->prepared; + return DAMACY_OK; +} + +static enum damacy_status +submit(struct damacy_executor* base, struct prepared_plan* plan, uint64_t id) +{ + struct test_executor* self = (void*)base; + self->last = plan; + ++self->submissions; + if (self->status != DAMACY_OK) + return self->status; + self->ids[self->accepted++] = id; + prepared_plan_destroy(plan); + return DAMACY_OK; +} + +static enum damacy_status +step(struct damacy_executor* base, int* changed) +{ + (void)base; + (void)changed; + return DAMACY_OK; +} + +static const struct damacy_planner_ops planner_ops = { .next = prepare }; +static const struct damacy_executor_ops executor_ops = { .submit = submit, + .step = step }; + +static int +test_bounded_preparation_and_retry(void) +{ + struct test_planner planner = { .base.ops = &planner_ops, .remaining = 4 }; + struct test_executor executor = { .base.ops = &executor_ops, + .status = DAMACY_AGAIN }; + struct prepared_plan* plans[2] = { 0 }; + struct damacy pipeline = { .planner = &planner.base, + .executor = &executor.base, + .plans = plans, + .queues.prepared_batches = 2 }; + stats_init(&pipeline.stats); + EXPECT(damacy_scheduler_step(&pipeline)); + EXPECT(planner.prepared == 2 && planner.remaining == 2); + EXPECT(pipeline.plan_count == 2 && pipeline.next_batch_id == 0); + struct prepared_plan* waiting = plans[0]; + EXPECT(executor.last == waiting); + EXPECT(!damacy_scheduler_step(&pipeline)); + EXPECT(executor.last == waiting && planner.prepared == 2); + executor.status = DAMACY_OK; + EXPECT(damacy_scheduler_step(&pipeline)); + EXPECT(pipeline.plan_count == 0 && executor.accepted == 2); + EXPECT(damacy_scheduler_step(&pipeline)); + EXPECT(pipeline.plan_count == 0 && executor.accepted == 4); + for (unsigned i = 0; i < 4; ++i) + EXPECT(executor.ids[i] == i); + EXPECT(pipeline.stats.plan.count == 4); + return 0; +} static int -test_plan_commit_releases_unsealed_slot(void) -{ - struct damacy d = { 0 }; - stats_init(&d.stats); - - struct damacy_batch_slot* slot = &d.batch_pool.slots[0]; - struct render_job* job = render_job_pool_for_batch_slot(&d.render_jobs, 0); - EXPECT(job); - - slot->state = BATCH_PLANNING; - slot->batch_id = 7; - slot->sample_seq_begin = 11; - slot->n_samples = 2; - - job->state = RENDER_JOB_READY; - job->batch_pool_slot = 0; - job->batch_id = 7; - job->n_chunks = 3; - - int changed = 0; - EXPECT(plan_commit(&d, 0, DAMACY_OK, 0.25f, &changed) == DAMACY_INVAL); - EXPECT(d.failed_status == DAMACY_INVAL); - EXPECT(changed == 1); - EXPECT(slot->state == BATCH_FREE); - EXPECT(slot->batch_id == 0); - EXPECT(slot->sample_seq_begin == 0); - EXPECT(slot->n_samples == 0); - EXPECT(job->state == RENDER_JOB_FREE); - EXPECT(job->batch_id == 0); - EXPECT(job->n_chunks == 0); - EXPECT(d.stats.plan.count == 1); +test_terminal_failures_preserve_ownership(void) +{ + struct test_planner planner = { .base.ops = &planner_ops, + .failure = DAMACY_NOTFOUND }; + struct test_executor executor = { .base.ops = &executor_ops, + .status = DAMACY_BUDGET }; + struct prepared_plan* plans[1] = { 0 }; + struct damacy pipeline = { .planner = &planner.base, + .executor = &executor.base, + .plans = plans, + .queues.prepared_batches = 1 }; + stats_init(&pipeline.stats); + EXPECT(damacy_scheduler_step(&pipeline)); + EXPECT(pipeline.failed_status == DAMACY_NOTFOUND && + executor.submissions == 0); + EXPECT(!damacy_scheduler_step(&pipeline)); + pipeline.failed_status = DAMACY_OK; + planner.failure = DAMACY_OK; + planner.remaining = 1; + EXPECT(damacy_scheduler_step(&pipeline)); + EXPECT(pipeline.failed_status == DAMACY_BUDGET); + EXPECT(pipeline.plan_count == 1 && plans[0] == executor.last); + EXPECT(pipeline.next_batch_id == 0 && executor.accepted == 0); + EXPECT(!damacy_scheduler_step(&pipeline)); + prepared_plan_destroy(plans[0]); return 0; } int main(void) { - RUN(test_plan_commit_releases_unsealed_slot); - printf("all damacy_plan tests passed\n"); + RUN(test_bounded_preparation_and_retry); + RUN(test_terminal_failures_preserve_ownership); return 0; } diff --git a/tests/test_group_chunks.c b/tests/test_group_chunks.c index 0a7aceac..8d3e6572 100644 --- a/tests/test_group_chunks.c +++ b/tests/test_group_chunks.c @@ -1,9 +1,9 @@ -// Unit tests for planner/group_chunks.c: stable counting-sort of +// Unit tests for executor/group_chunks.c: stable counting-sort of // chunk_plans by read_op_idx. Synthetic inputs only. +#include "executor/dispatch.h" +#include "executor/group_chunks.h" #include "expect.h" -#include "planner/group_chunks.h" -#include "planner/planner.h" #include #include @@ -26,10 +26,10 @@ mk_cp(struct chunk_plan* cp, } static enum damacy_status -run_group(struct planner_output* out) +run_group(struct dispatch_output* out) { - uint32_t* u32 = (uint32_t*)calloc((size_t)out->n_read_ops + 1u, - sizeof(uint32_t)); + uint32_t* u32 = + (uint32_t*)calloc((size_t)out->n_read_ops + 1u, sizeof(uint32_t)); struct chunk_plan* tmp = (struct chunk_plan*)calloc( out->n_chunk_plans ? out->n_chunk_plans : 1, sizeof(struct chunk_plan)); if (!out->read_op_groups && out->n_read_ops > 0) { @@ -47,7 +47,7 @@ run_group(struct planner_output* out) static int test_empty(void) { - struct planner_output out = { 0 }; + struct dispatch_output out = { 0 }; struct chunk_plan c = { 0 }; out.chunk_plans = &c; out.chunk_plans_cap = 1; @@ -62,7 +62,7 @@ test_single_chunk(void) { struct chunk_plan chunks[1]; mk_cp(&chunks[0], 0, 100, 0, 42); - struct planner_output out = { + struct dispatch_output out = { .chunk_plans = chunks, .chunk_plans_cap = 1, .n_chunk_plans = 1, @@ -83,7 +83,7 @@ test_already_sorted(void) mk_cp(&chunks[1], 0, 4096, 0, 101); mk_cp(&chunks[2], 1, 0, 0, 200); mk_cp(&chunks[3], 1, 4096, 0, 201); - struct planner_output out = { + struct dispatch_output out = { .chunk_plans = chunks, .chunk_plans_cap = 4, .n_chunk_plans = 4, @@ -111,7 +111,7 @@ test_interleaved_groups(void) mk_cp(&chunks[3], 2, 0, 2, 300); mk_cp(&chunks[4], 1, 100, 1, 201); mk_cp(&chunks[5], 0, 200, 0, 102); - struct planner_output out = { + struct dispatch_output out = { .chunk_plans = chunks, .chunk_plans_cap = 6, .n_chunk_plans = 6, @@ -139,7 +139,7 @@ test_sparse_read_op(void) mk_cp(&chunks[0], 2, 0, 0, 1); mk_cp(&chunks[1], 2, 100, 0, 2); mk_cp(&chunks[2], 2, 200, 0, 3); - struct planner_output out = { + struct dispatch_output out = { .chunk_plans = chunks, .chunk_plans_cap = 3, .n_chunk_plans = 3, @@ -159,7 +159,7 @@ test_invalid_read_op_idx(void) struct chunk_plan chunks[2]; mk_cp(&chunks[0], 0, 0, 0, 1); mk_cp(&chunks[1], 5, 0, 0, 2); // out of range - struct planner_output out = { + struct dispatch_output out = { .chunk_plans = chunks, .chunk_plans_cap = 2, .n_chunk_plans = 2, @@ -180,7 +180,7 @@ test_sample_distribution_preserved(void) mk_cp(&chunks[2], 1, 100, /*sample*/ 7, 3); mk_cp(&chunks[3], 0, 100, /*sample*/ 4, 4); mk_cp(&chunks[4], 0, 200, /*sample*/ 0, 5); - struct planner_output out = { + struct dispatch_output out = { .chunk_plans = chunks, .chunk_plans_cap = 5, .n_chunk_plans = 5, @@ -221,7 +221,7 @@ test_groups_emitted(void) chunks[3].decompressed_nbytes = 40; chunks[4].decompressed_nbytes = 50; chunks[5].decompressed_nbytes = 60; - struct planner_output out = { + struct dispatch_output out = { .chunk_plans = chunks, .chunk_plans_cap = 6, .n_chunk_plans = 6, @@ -257,7 +257,7 @@ test_groups_skip_sparse(void) chunks[0].decompressed_nbytes = 7; chunks[1].decompressed_nbytes = 7; chunks[2].decompressed_nbytes = 7; - struct planner_output out = { + struct dispatch_output out = { .chunk_plans = chunks, .chunk_plans_cap = 3, .n_chunk_plans = 3, @@ -278,9 +278,18 @@ static int test_iterator_walk(void) { struct read_op_group groups[3] = { - { .read_op_idx = 0, .first_chunk = 0, .n_chunks = 2, .total_decompressed = 10 }, - { .read_op_idx = 1, .first_chunk = 2, .n_chunks = 1, .total_decompressed = 20 }, - { .read_op_idx = 2, .first_chunk = 3, .n_chunks = 4, .total_decompressed = 30 }, + { .read_op_idx = 0, + .first_chunk = 0, + .n_chunks = 2, + .total_decompressed = 10 }, + { .read_op_idx = 1, + .first_chunk = 2, + .n_chunks = 1, + .total_decompressed = 20 }, + { .read_op_idx = 2, + .first_chunk = 3, + .n_chunks = 4, + .total_decompressed = 30 }, }; struct read_op_group_iterator it; read_op_group_iterator_init(&it, groups, 3, 1); diff --git a/tests/test_planner.c b/tests/test_planner.c index 63a98759..5617b58b 100644 --- a/tests/test_planner.c +++ b/tests/test_planner.c @@ -378,7 +378,7 @@ test_single_chunk_aligned(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -453,7 +453,7 @@ test_multi_chunk_partial(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -520,7 +520,7 @@ test_two_samples_indices(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -568,7 +568,7 @@ test_empty_chunk_becomes_fill(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -620,7 +620,7 @@ test_fill_value_int16_neg1(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -663,7 +663,7 @@ test_fill_value_f32_nan(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -708,7 +708,7 @@ test_missing_shard_becomes_fill(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -749,7 +749,7 @@ test_page_alignment(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -806,7 +806,7 @@ run_blosc_codec_id_case(const char* cname, uint8_t expected_codec_id) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -848,7 +848,7 @@ run_blosc_lz4_rejected_case(const char* cname) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -901,7 +901,7 @@ test_codec_id_none(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -980,7 +980,7 @@ test_unsupported_source_dtype(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -1082,7 +1082,7 @@ test_unsharded_single_chunk(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -1131,7 +1131,7 @@ test_unsharded_multi_chunk(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -1203,7 +1203,7 @@ test_sharded_index_start(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -1248,7 +1248,7 @@ test_coalesce_adjacent_pages(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -1295,7 +1295,7 @@ test_coalesce_gap_blocks_fusion(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -1343,7 +1343,7 @@ test_coalesce_fill_does_not_block_fusion(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -1405,7 +1405,7 @@ test_coalesce_non_monotonic_shard(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, @@ -1452,7 +1452,7 @@ test_coalesce_cross_sample(void) struct read_op reads[8] = { 0 }; struct chunk_plan chunks[8] = { 0 }; struct sample_plan samples[4] = { 0 }; - struct planner_output out = { + struct dispatch_output out = { .read_ops = reads, .read_ops_cap = 8, .chunk_plans = chunks, diff --git a/tests/test_render_job.c b/tests/test_render_job.c index 4b7b28e7..45cb8e43 100644 --- a/tests/test_render_job.c +++ b/tests/test_render_job.c @@ -21,7 +21,7 @@ wire_job_storage(struct render_job* job, } static int -test_planner_output_borrows_job_storage(void) +test_dispatch_output_borrows_job_storage(void) { struct render_job job; struct read_op reads[2]; @@ -30,7 +30,7 @@ test_planner_output_borrows_job_storage(void) struct read_op_group groups[2]; wire_job_storage(&job, reads, chunks, samples, groups); - struct planner_output out = render_job_planner_output(&job, 2); + struct dispatch_output out = render_job_dispatch_output(&job, 2); EXPECT(out.read_ops == reads); EXPECT(out.chunk_plans == chunks); EXPECT(out.sample_plans == samples); @@ -45,7 +45,7 @@ test_commit_and_find_oldest_work(void) { struct render_job_pool pool; memset(&pool, 0, sizeof(pool)); - struct planner_output out = { + struct dispatch_output out = { .n_chunk_plans = 3, .n_chunks_to_load = 2, .n_loads_issued = 1, @@ -75,7 +75,7 @@ test_commit_and_find_oldest_work(void) int main(void) { - RUN(test_planner_output_borrows_job_storage); + RUN(test_dispatch_output_borrows_job_storage); RUN(test_commit_and_find_oldest_work); printf("all render_job tests passed\n"); return 0; diff --git a/tests/test_scheduler.c b/tests/test_scheduler.c index 75abda94..a111f553 100644 --- a/tests/test_scheduler.c +++ b/tests/test_scheduler.c @@ -31,7 +31,7 @@ static int test_create_destroy(void) { struct ctx c = { 0, 0 }; - struct scheduler* s = scheduler_create(step_count, &c, 1000000, NULL); + struct scheduler* s = scheduler_create(step_count, &c, 1000000, NULL, NULL); EXPECT(s != NULL); scheduler_destroy(s); scheduler_destroy(NULL); // NULL-safe @@ -43,7 +43,7 @@ test_step_runs_periodically(void) { struct ctx c = { 0, -1 }; // never signal; just count struct scheduler* s = - scheduler_create(step_count, &c, 500000, NULL); // 500 µs + scheduler_create(step_count, &c, 500000, NULL, NULL); // 500 µs EXPECT(s != NULL); platform_sleep_ns(20000000); // 20 ms → expect >> 10 ticks scheduler_lock(s); @@ -58,7 +58,7 @@ static int test_signal_wakes_waiter(void) { struct ctx c = { 0, 5 }; - struct scheduler* s = scheduler_create(step_count, &c, 500000, NULL); + struct scheduler* s = scheduler_create(step_count, &c, 500000, NULL, NULL); EXPECT(s != NULL); scheduler_lock(s); while (c.n_steps < c.target) @@ -99,7 +99,7 @@ static int test_external_broadcast(void) { struct scheduler* s = - scheduler_create(step_noop, NULL, 1000000000, NULL); // 1 s tick + scheduler_create(step_noop, NULL, 1000000000, NULL, NULL); // 1 s tick EXPECT(s != NULL); int flag = 0; struct broadcaster_args a = { s, &flag }; @@ -119,9 +119,9 @@ test_external_broadcast(void) static int test_invalid_args(void) { - EXPECT(scheduler_create(NULL, NULL, 1000, NULL) == NULL); - EXPECT(scheduler_create(step_noop, NULL, 0, NULL) == NULL); - EXPECT(scheduler_create(step_noop, NULL, -1, NULL) == NULL); + EXPECT(scheduler_create(NULL, NULL, 1000, NULL, NULL) == NULL); + EXPECT(scheduler_create(step_noop, NULL, 0, NULL, NULL) == NULL); + EXPECT(scheduler_create(step_noop, NULL, -1, NULL, NULL) == NULL); return 0; } diff --git a/tests/test_wave_pool.c b/tests/test_wave_pool.c index 754c708b..43114189 100644 --- a/tests/test_wave_pool.c +++ b/tests/test_wave_pool.c @@ -7,8 +7,8 @@ #include "damacy.h" #include "damacy_limits.h" #include "damacy_log.h" +#include "executor/dispatch.h" #include "expect.h" -#include "planner/planner.h" #include "render_job/render_job.h" #include "store/store.h" #include "wave/input_slot.h" diff --git a/tests/write_zarr.py b/tests/write_zarr.py index 16eee709..d7fa1112 100644 --- a/tests/write_zarr.py +++ b/tests/write_zarr.py @@ -28,7 +28,7 @@ def parse_shape(s: str) -> tuple[int, ...]: return tuple(int(x) for x in s.split(",")) -def make_compressors(codec: str, dtype: np.dtype): +def make_compressors(codec: str, dtype: np.dtype, shuffle: str): if codec == "none": return [] if codec == "zstd": @@ -47,7 +47,7 @@ def make_compressors(codec: str, dtype: np.dtype): BloscCodec( cname=BloscCname.zstd, clevel=clevel, - shuffle=BloscShuffle.shuffle, + shuffle=BloscShuffle[shuffle], typesize=int(dtype.itemsize), ) ] @@ -82,6 +82,9 @@ def main() -> int: default="zstd", help="inner codec: none | zstd | blosc-zstd | blosc-zstd-l", ) + ap.add_argument( + "--shuffle", choices=["noshuffle", "shuffle", "bitshuffle"], default="shuffle" + ) args = ap.parse_args() if not (len(args.shape) == len(args.inner) == len(args.shard)): @@ -109,7 +112,7 @@ def main() -> int: dtype=np_dtype, chunks=args.inner, shards=args.shard, - compressors=make_compressors(args.codec, np_dtype), + compressors=make_compressors(args.codec, np_dtype, args.shuffle), ) arr[...] = data return 0 From 70ee773cc0a9c7c6d78cd6635904c4aa873b715a Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Mon, 14 Sep 2026 19:35:05 +0000 Subject: [PATCH 2/7] ci: fix TSan address layout --- .github/workflows/tsan.yml | 1 + Dockerfile | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/tsan.yml b/.github/workflows/tsan.yml index 6bce38ae..028da67c 100644 --- a/.github/workflows/tsan.yml +++ b/.github/workflows/tsan.yml @@ -46,4 +46,5 @@ jobs: --security-opt seccomp=unconfined \ -e TSAN_OPTIONS="suppressions=/workspace/damacy/tests/tsan-suppressions.txt halt_on_error=1 second_deadlock_stack=1" \ "$IMAGE" \ + setarch --addr-no-randomize \ ctest --test-dir build -L tsan --output-on-failure diff --git a/Dockerfile b/Dockerfile index d16ca40c..2f881af6 100644 --- a/Dockerfile +++ b/Dockerfile @@ -45,6 +45,7 @@ RUN apt-get update \ ca-certificates \ curl \ xz-utils \ + util-linux \ cmake \ ninja-build \ pkg-config \ From 557dff54e9807571e0cf8c9d30f0aa85e8a06ba2 Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Mon, 14 Sep 2026 19:39:52 +0000 Subject: [PATCH 3/7] test: use string Blosc options --- tests/write_zarr.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/write_zarr.py b/tests/write_zarr.py index d7fa1112..bc2e7238 100644 --- a/tests/write_zarr.py +++ b/tests/write_zarr.py @@ -19,7 +19,7 @@ import sys import numpy as np -from zarr.codecs import BloscCname, BloscCodec, BloscShuffle, ZstdCodec +from zarr.codecs import BloscCodec, ZstdCodec import zarr @@ -45,9 +45,9 @@ def make_compressors(codec: str, dtype: np.dtype, shuffle: str): raise SystemExit(f"unknown --codec {codec!r}") return [ BloscCodec( - cname=BloscCname.zstd, + cname="zstd", clevel=clevel, - shuffle=BloscShuffle[shuffle], + shuffle=shuffle, typesize=int(dtype.itemsize), ) ] From 5126f0958c0be3922a0a00c49d127ba62ddbc6a9 Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Mon, 14 Sep 2026 21:08:39 +0000 Subject: [PATCH 4/7] docs: qualify CUDA benchmark results --- dev/cpu-pipeline-validation.md | 22 ++++++++++++---------- dev/cpu-pipeline.md | 7 ++++--- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/dev/cpu-pipeline-validation.md b/dev/cpu-pipeline-validation.md index 27ec0e2e..ff434118 100644 --- a/dev/cpu-pipeline-validation.md +++ b/dev/cpu-pipeline-validation.md @@ -52,7 +52,7 @@ from 2.016 GiB with one worker to 2.255 GiB with sixteen, including conservative codec workspace allowances. That reservation excludes metadata, plans, reader queues, thread stacks, and allocator overhead; it is not a process RSS limit. -## CUDA comparison and remaining performance issue +## CUDA comparison and uncertainty | Pair | Baseline GB/s | Refactor GB/s | | ---: | ---: | ---: | @@ -61,21 +61,23 @@ queues, thread stacks, and allocator overhead; it is not a process RSS limit. | 3 | 12.409 | 11.646 | | Median | 12.409 | 11.684 | -Median useful throughput is **5.8% lower** after the refactor. This is an -open performance regression on this workload; CUDA throughput parity has not -been established. The lower first baseline result is retained in the table. -Three runs do not characterize all filesystem or scheduling variability. +Median useful throughput is **5.8% lower** in the refactor runs. Baseline +throughput ranges from 11.834 to 12.516 GB/s, a spread of about 6%. With only +three pairs and uncontrolled shared NFS traffic, these results do not establish +whether the difference comes from code changes or storage variability. +Performance parity has also not been established. The measured decode kernels take about 2.515 seconds in both versions. The additional time appears between decode waves, with substantial gaps at batch boundaries. A separate timing trace recorded approximately 45 ms of shared planning and 44 ms of CUDA dispatch preparation across 35 batches including -warmup. These measurements narrow the follow-up to host preparation and -scheduling; they do not establish one root cause. +warmup. These timings do not identify how much of the throughput difference +comes from storage waits, host preparation, or scheduling. -Correctness and retained-result lifetimes pass on CUDA. Further work on CUDA -preparation and scheduling should preserve the owned-plan boundary and be -measured separately from the future query features. +Correctness and retained-result lifetimes pass on CUDA. Before attributing a +performance difference to code, a follow-up should randomize the order within +pairs and collect more repetitions under controlled storage conditions. A +comparison using local storage could help isolate NFS effects. ## Checks diff --git a/dev/cpu-pipeline.md b/dev/cpu-pipeline.md index 15fc73aa..3a76bec4 100644 --- a/dev/cpu-pipeline.md +++ b/dev/cpu-pipeline.md @@ -159,9 +159,10 @@ work items. ## Validation Measured results and build evidence are recorded in -[CPU pipeline validation](cpu-pipeline-validation.md). The current L40 -comparison shows about 6% lower CUDA throughput; further scheduling work is -needed before claiming performance parity. +[CPU pipeline validation](cpu-pipeline-validation.md). The L40 comparison +shows 5.8% lower median CUDA throughput across three pairs on shared NFS. +The small sample and uncontrolled storage traffic do not establish a code +regression; a controlled comparison is needed to attribute the difference. The CPU milestone checks independent crop values across codecs and source dtypes, missing fills, bfloat16 conversion, duplicate chunk use, corrupt input, From 654c27178485003fd48e7b16f6d088f3f877d056 Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Thu, 24 Sep 2026 21:16:47 +0000 Subject: [PATCH 5/7] fix: free CPU plans after reads finish --- src/executor/cpu_executor.c | 2 + tests/CMakeLists.txt | 1 + tests/test_cpu_executor.c | 275 ++++++++++++++++++++++++++++++++++++ 3 files changed, 278 insertions(+) create mode 100644 tests/test_cpu_executor.c diff --git a/src/executor/cpu_executor.c b/src/executor/cpu_executor.c index dc5b735f..e44df4b9 100644 --- a/src/executor/cpu_executor.c +++ b/src/executor/cpu_executor.c @@ -318,6 +318,8 @@ cpu_stop(struct damacy_executor* base) free(wave->assemble_ms); free(wave->output_bytes); *wave = (struct cpu_wave){ 0 }; + } + for (unsigned i = 0; i < 2; ++i) { prepared_plan_destroy(self->slots[i].plan); buffer_release(self->slots[i].buffer); self->slots[i] = (struct cpu_slot){ 0 }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index bce68503..f4b0c6fe 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -64,6 +64,7 @@ add_damacy_test(test_scheduler scheduler platform log) if(NOT DAMACY_FUZZ) add_damacy_test(test_damacy_plan damacy) add_damacy_test(test_cpu_pipeline damacy test_fixture Threads::Threads) + add_damacy_test(test_cpu_executor cpu_executor) set_tests_properties(test_cpu_pipeline PROPERTIES TIMEOUT 120) endif() diff --git a/tests/test_cpu_executor.c b/tests/test_cpu_executor.c new file mode 100644 index 00000000..01ddff14 --- /dev/null +++ b/tests/test_cpu_executor.c @@ -0,0 +1,275 @@ +#include "damacy_stats.h" +#include "expect.h" +#include "pipeline/components.h" +#include "store/store_internal.h" + +#include +#include + +enum +{ + MAX_CHUNKS = 32, + MAX_READS = 64, + SHARD_BYTES = 128, + CHUNK_BYTES = 8, +}; + +struct read_record +{ + unsigned shard; + uint64_t offset; + size_t bytes; +}; + +struct pending_reads +{ + struct store_read reads[MAX_CHUNKS]; + size_t count; + int active; +}; + +struct test_store +{ + struct store base; + uint16_t data[3][SHARD_BYTES / sizeof(uint16_t)]; + struct read_record records[MAX_READS]; + struct pending_reads events[MAX_READS]; + unsigned n_reads; + unsigned n_events; + unsigned pending; + unsigned capacity; + unsigned rejected; + unsigned waited; + unsigned completion_errors; + uint64_t held_event; + enum damacy_status submit_status; + enum damacy_status read_status; +}; + +static enum damacy_status +complete_reads(struct test_store* store, struct store_event event) +{ + if (!event.impl) + return DAMACY_OK; + struct pending_reads* pending = event.impl; + enum damacy_status status = store->read_status; + for (size_t i = 0; i < pending->count && status == DAMACY_OK; ++i) { + const struct store_read* read = &pending->reads[i]; + unsigned shard = (unsigned)(read->key[0] - 'a'); + if (shard >= 3 || read->offset > SHARD_BYTES || + read->len > SHARD_BYTES - read->offset) { + status = DAMACY_IO; + break; + } + memcpy(read->dst, (char*)store->data[shard] + read->offset, read->len); + } + store->pending -= (unsigned)pending->count; + pending->active = 0; + if (status != DAMACY_OK) + ++store->completion_errors; + return status; +} + +static struct store_submit_result +submit_reads(struct store* base, const struct store_read* reads, size_t count) +{ + struct test_store* store = (void*)base; + if (store->submit_status != DAMACY_OK) + return (struct store_submit_result){ .status = store->submit_status }; + if (!count) + return (struct store_submit_result){ .status = DAMACY_OK }; + if (count > store->capacity - store->pending) { + ++store->rejected; + return (struct store_submit_result){ .status = DAMACY_AGAIN }; + } + if (count > MAX_CHUNKS || count > MAX_READS - store->n_reads || + store->n_events == MAX_READS) + return (struct store_submit_result){ .status = DAMACY_IO }; + struct pending_reads* pending = &store->events[store->n_events++]; + pending->count = count; + pending->active = 1; + memcpy(pending->reads, reads, count * sizeof(*reads)); + for (size_t i = 0; i < count; ++i) + store->records[store->n_reads++] = + (struct read_record){ .shard = (unsigned)(reads[i].key[0] - 'a'), + .offset = reads[i].offset, + .bytes = reads[i].len }; + store->pending += (unsigned)count; + return (struct store_submit_result){ + .status = DAMACY_OK, .event = { .seq = store->n_events, .impl = pending } + }; +} + +static struct store_event_poll +query_reads(struct store* base, struct store_event event) +{ + struct test_store* store = (void*)base; + if (event.impl && event.seq == store->held_event) + return (struct store_event_poll){ .status = DAMACY_OK }; + return (struct store_event_poll){ .status = complete_reads(store, event), + .ready = 1 }; +} + +static enum damacy_status +wait_reads(struct store* base, struct store_event event) +{ + struct test_store* store = (void*)base; + ++store->waited; + return complete_reads(store, event); +} + +static const struct store_vtable test_store_ops = { .submit = submit_reads, + .event_query = query_reads, + .event_wait = wait_reads }; + +static void +store_init(struct test_store* store, unsigned capacity) +{ + *store = (struct test_store){ .base = { .vt = &test_store_ops }, + .capacity = capacity }; + for (unsigned shard = 0; shard < 3; ++shard) + for (unsigned i = 0; i < SHARD_BYTES / sizeof(uint16_t); ++i) + store->data[shard][i] = (uint16_t)(100 * shard + i); +} + +struct test_chunk +{ + unsigned shard; + uint64_t offset; + uint8_t missing; +}; + +struct plan_storage +{ + struct plan_array array; + struct plan_chunk chunks[MAX_CHUNKS]; + struct plan_region regions[2]; + struct plan_use uses[2 * MAX_CHUNKS]; + char paths[MAX_CHUNKS][2]; +}; + +static struct damacy_batch_spec +output_spec(uint32_t count) +{ + return (struct damacy_batch_spec){ .dtype = DAMACY_F32, + .sample_shape = { count * 4 }, + .sample_rank = 1, + .samples_per_batch = 2 }; +} + +static struct prepared_plan* +make_plan(const struct test_chunk* chunks, uint32_t count) +{ + if (!count || count > MAX_CHUNKS) + return NULL; + struct prepared_plan* plan = calloc(1, sizeof(*plan)); + struct plan_storage* storage = calloc(1, sizeof(*storage)); + if (!plan || !storage) { + free(plan); + free(storage); + return NULL; + } + *plan = + (struct prepared_plan){ .output = output_spec(count), + .arrays = &storage->array, + .chunks = storage->chunks, + .regions = storage->regions, + .uses = storage->uses, + .n_arrays = 1, + .n_chunks = count, + .n_regions = 2, + .n_uses = 2 * count, + .allocated_bytes = sizeof(*plan) + sizeof(*storage), + .storage = storage }; + storage->array.metadata = + (struct zarr_metadata){ .rank = 1, + .dtype = dtype_u16, + .shape = { count * 4 }, + .inner_chunk_shape = { 4 }, + .inner_codec = { .id = CODEC_NONE } }; + uint16_t fill = 999; + memcpy(storage->array.metadata.fill_value, &fill, sizeof(fill)); + for (unsigned sample = 0; sample < 2; ++sample) + storage->regions[sample] = (struct plan_region){ + .sample = sample, .source = { .rank = 1, .dims = { { 0, count * 4 } } } + }; + for (uint32_t i = 0; i < count; ++i) { + storage->paths[i][0] = (char)('a' + chunks[i].shard); + storage->chunks[i] = + (struct plan_chunk){ .path = storage->paths[i], + .offset = chunks[i].offset, + .coordinate = { i }, + .encoded_bytes = chunks[i].missing ? 0 : CHUNK_BYTES, + .decoded_bytes = CHUNK_BYTES, + .first_use = 2 * i, + .missing = chunks[i].missing }; + storage->uses[2 * i] = + (struct plan_use){ .chunk = i, .region = 0, .next = 2 * i + 1 }; + storage->uses[2 * i + 1] = + (struct plan_use){ .chunk = i, .region = 1, .next = UINT32_MAX }; + } + return plan; +} + +static int +start_executor(struct damacy_reader* reader, + uint32_t workers, + uint32_t count, + uint64_t budget, + struct damacy_stats* stats, + struct damacy_executor** out) +{ + struct damacy_cpu_config config = { .decode_workers = workers, + .max_encoded_chunk_bytes = CHUNK_BYTES, + .max_decoded_chunk_bytes = CHUNK_BYTES, + .max_memory_bytes = budget }; + EXPECT(damacy_cpu_executor_create(reader, &config, out) == DAMACY_OK); + struct damacy_batch_spec output = output_spec(count); + EXPECT((*out)->ops->start(*out, &output, stats) == DAMACY_OK); + return 0; +} + +static int +test_read_errors_and_shutdown(void) +{ + const struct test_chunk chunks[] = { + { .offset = 0 }, { .offset = 8 }, { .offset = 16 }, { .offset = 24 } + }; + for (unsigned failure = 0; failure < 3; ++failure) { + struct test_store store; + store_init(&store, 4); + if (failure == 0) + store.submit_status = DAMACY_IO; + else if (failure == 1) + store.read_status = DAMACY_IO; + else + store.held_event = 2; + struct damacy_reader reader = { .store = &store.base, + .max_inflight_reads = 2 }; + struct damacy_stats stats = { 0 }; + struct damacy_executor* executor = NULL; + EXPECT(start_executor(&reader, 2, 4, 8 << 20, &stats, &executor) == 0); + struct prepared_plan* plan = make_plan(chunks, 4); + EXPECT(plan); + EXPECT(executor->ops->submit(executor, plan, 0) == DAMACY_OK); + int changed = 0; + EXPECT(executor->ops->step(executor, &changed) == + (failure == 2 ? DAMACY_OK : DAMACY_IO)); + struct damacy_batch* batch = NULL; + EXPECT(executor->ops->take(executor, &batch) == DAMACY_AGAIN); + damacy_executor_destroy(executor); + EXPECT(store.pending == 0); + if (failure) + EXPECT(store.waited == 1); + if (failure == 2) + EXPECT(store.completion_errors == 0); + } + return 0; +} + +int +main(void) +{ + RUN(test_read_errors_and_shutdown); + return 0; +} From c78d66a8673c4e7ebdb40eeb85b4424ee8f7e13c Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Thu, 24 Sep 2026 21:18:52 +0000 Subject: [PATCH 6/7] feat: let planners share metadata --- dev/cpu-pipeline.md | 7 ++--- docs/pipeline.md | 10 ++++--- python/damacy/__init__.py | 16 ++++++++--- python/tests/test_components.py | 31 ++++++++++++++++++---- src/damacy_pipeline.h | 4 ++- src/pipeline/components.c | 10 +------ src/pipeline/components.h | 4 +-- src/pipeline/zarr_planner.c | 26 +++--------------- tests/test_cpu_pipeline.c | 47 +++++++++++++++++++++++++++++++++ 9 files changed, 104 insertions(+), 51 deletions(-) diff --git a/dev/cpu-pipeline.md b/dev/cpu-pipeline.md index 3a76bec4..8abd6b8d 100644 --- a/dev/cpu-pipeline.md +++ b/dev/cpu-pipeline.md @@ -97,9 +97,10 @@ Views survive batch release and pipeline shutdown. Explicit CUDA devices retain the primary context and completion stream until the last buffer is released. Caller-owned CUDA contexts must outlive their views. -Planners, executors, metadata providers, and metadata readers reject use by two -active pipelines. They can be reused after shutdown. Python retains dependencies; -C borrows them until shutdown and requires reverse-order destruction. Queue and +Planners and executors reject use by two active pipelines. They can be reused +after shutdown. Metadata providers and metadata readers hold only settings that +planners copy, so planners may share them. Python retains dependencies; C +borrows them until shutdown and requires reverse-order destruction. Queue and buffer saturation report retriable backpressure, not a storage error. Closing a pipeline stops preparation, wakes blocked pops, joins execution, and releases queued work before its dependencies can disappear. diff --git a/docs/pipeline.md b/docs/pipeline.md index 86183e9e..20080674 100644 --- a/docs/pipeline.md +++ b/docs/pipeline.md @@ -146,10 +146,12 @@ requires closing the pipeline and constructing another one. ## Lifetimes and interoperation Use a context manager or call `Pipeline.close()`. Python retains the injected -components and their dependencies. Planners, executors, metadata providers, -and metadata readers each serve one active pipeline; simultaneous reuse is -rejected. They may be reused after close. Closing stops pending work and wakes -blocked pops. Reuse creates fresh active caches and execution resources. +components and their dependencies. Planners and executors each serve one +active pipeline; simultaneous reuse is rejected. They may be reused after +close. Metadata providers and metadata readers hold only settings. Each planner +copies them and builds its own caches, so any number of planners may share one. +Closing stops pending work and wakes blocked pops. Reuse creates fresh active +caches and execution resources. `np.from_dlpack(batch)` produces a CPU view; `torch.from_dlpack(batch)` accepts CPU or CUDA results. A DLPack view remains valid after releasing the `Batch` and diff --git a/python/damacy/__init__.py b/python/damacy/__init__.py index 848992f5..e8c2de09 100644 --- a/python/damacy/__init__.py +++ b/python/damacy/__init__.py @@ -881,7 +881,10 @@ def __init__(self, *, workers: int = 8, max_inflight_reads: int = 4096) -> None: class FileMetadataReader: - """Configure a separate asynchronous queue for filesystem metadata reads.""" + """Configure a separate asynchronous queue for filesystem metadata reads. + + Holds settings only; metadata providers copy them, so it can be shared. + """ __slots__ = ("_native",) @@ -903,7 +906,11 @@ def __init__( class ZarrMetadata: - """Provide Zarr v3 array metadata and shard indexes for sample URIs.""" + """Provide Zarr v3 array metadata and shard indexes for sample URIs. + + Holds settings only. Each planner copies them and keeps its own caches, so + planners can share one instance. + """ __slots__ = ("_native", "cache", "reader") @@ -1381,8 +1388,9 @@ class Pipeline: ``planner`` resolves source metadata into owned chunk plans. ``executor`` reads, decodes, and assembles them. ``output`` defines the batch tensor; - ``queues`` bounds preparation. Components serve one active pipeline and - may be reused after it closes. Exported tensors retain their storage. + ``queues`` bounds preparation. The planner and executor serve one active + pipeline and may be reused after it closes. Exported tensors retain their + storage. ``Pipeline(Config(...))`` composes the CUDA pipeline for existing callers. For CUDA, pass an explicit executor device or make a CUDA context current diff --git a/python/tests/test_components.py b/python/tests/test_components.py index 2877f028..262f9bb8 100644 --- a/python/tests/test_components.py +++ b/python/tests/test_components.py @@ -15,12 +15,16 @@ import pytest -def planner(**limits): +def zarr_metadata(): + return damacy.ZarrMetadata( + reader=damacy.FileMetadataReader(concurrency=2), + cache=damacy.MetadataCache(array_entries=16, shard_index_entries=64), + ) + + +def planner(metadata=None, **limits): return damacy.ChunkPlanner( - metadata=damacy.ZarrMetadata( - reader=damacy.FileMetadataReader(concurrency=2), - cache=damacy.MetadataCache(array_entries=16, shard_index_entries=64), - ), + metadata=metadata or zarr_metadata(), limits=damacy.PlanLimits( **( {"max_chunks": 128, "max_chunk_bytes": 1024, "max_shards_per_sample": 4} @@ -180,6 +184,23 @@ def test_components_exclusive_and_reusable(tiny_zarr): ) +def test_planners_share_metadata(tiny_zarr): + metadata = zarr_metadata() + pipelines = [pipeline(planner=planner(metadata)) for _ in range(2)] + try: + for p in pipelines: + p.push([sample(tiny_zarr)]) + for p in pipelines: + with p.pop() as batch: + np.testing.assert_array_equal( + np.from_dlpack(batch), + np.arange(128, dtype=np.float32).reshape(1, 8, 16), + ) + finally: + for p in pipelines: + p.close() + + def test_mixed_source_types(tmp_path): a = np.arange(16, dtype=np.uint16).reshape(4, 4) b = (np.arange(16, dtype=np.float32) / 8 - 2).reshape(4, 4) diff --git a/src/damacy_pipeline.h b/src/damacy_pipeline.h index 2aed9149..bab1b71f 100644 --- a/src/damacy_pipeline.h +++ b/src/damacy_pipeline.h @@ -99,7 +99,9 @@ extern "C" void damacy_executor_destroy(struct damacy_executor* executor); // Components are borrowed until shutdown; each planner/executor may serve - // one active pipeline. Readers and metadata outlive their dependents. + // one active pipeline. A reader must outlive its executors. Metadata and + // metadata readers hold only settings, copied by the objects created from + // them, so they can be shared or destroyed at any time. enum damacy_status damacy_pipeline_create( struct damacy_planner* planner, struct damacy_executor* executor, diff --git a/src/pipeline/components.c b/src/pipeline/components.c index aeffd781..55ece990 100644 --- a/src/pipeline/components.c +++ b/src/pipeline/components.c @@ -96,10 +96,6 @@ damacy_file_metadata_reader_create(uint32_t concurrency, void damacy_metadata_reader_destroy(struct damacy_metadata_reader* reader) { - if (reader && reader->active) { - log_error("metadata reader is still in use"); - return; - } free(reader); } @@ -116,7 +112,7 @@ damacy_zarr_metadata_create(struct damacy_metadata_reader* reader, struct damacy_metadata* metadata = calloc(1, sizeof(*metadata)); if (!metadata) return DAMACY_OOM; - metadata->reader = reader; + metadata->reader = *reader; metadata->cache = *cache; *out = metadata; return DAMACY_OK; @@ -125,10 +121,6 @@ damacy_zarr_metadata_create(struct damacy_metadata_reader* reader, void damacy_metadata_destroy(struct damacy_metadata* metadata) { - if (metadata && metadata->active) { - log_error("metadata provider is still in use"); - return; - } free(metadata); } diff --git a/src/pipeline/components.h b/src/pipeline/components.h index 5dd3992b..6838f3aa 100644 --- a/src/pipeline/components.h +++ b/src/pipeline/components.h @@ -17,14 +17,12 @@ struct damacy_metadata_reader { uint32_t concurrency; struct damacy_latency_model latency; - _Atomic int active; }; struct damacy_metadata { - struct damacy_metadata_reader* reader; + struct damacy_metadata_reader reader; struct damacy_metadata_cache_config cache; - _Atomic int active; }; struct damacy_planner_ops diff --git a/src/pipeline/zarr_planner.c b/src/pipeline/zarr_planner.c index edf45c86..840c5bfb 100644 --- a/src/pipeline/zarr_planner.c +++ b/src/pipeline/zarr_planner.c @@ -13,7 +13,7 @@ struct zarr_planner { struct damacy_planner base; - struct damacy_metadata* metadata; + struct damacy_metadata metadata; struct damacy_plan_limits limits; struct damacy_batch_spec output; struct damacy_queue_limits queues; @@ -29,7 +29,6 @@ struct zarr_planner uint64_t pushed; uint64_t planned; uint64_t watermark; - int bound; }; static void @@ -60,11 +59,6 @@ zarr_stop(struct damacy_planner* base) clear_samples(self); free(self->samples); self->samples = NULL; - if (self->bound) { - self->metadata->active = 0; - self->metadata->reader->active = 0; - self->bound = 0; - } } static enum damacy_status @@ -73,9 +67,7 @@ zarr_start(struct damacy_planner* base, const struct damacy_queue_limits* queues) { struct zarr_planner* self = (void*)base; - struct damacy_metadata* metadata = self->metadata; - if (metadata->active || metadata->reader->active) - return DAMACY_INVAL; + const struct damacy_metadata* metadata = &self->metadata; uint64_t floor = (uint64_t)queues->lookahead_samples + output->samples_per_batch; if (metadata->cache.array_entries < floor || @@ -90,18 +82,8 @@ zarr_start(struct damacy_planner* base, self->output = *output; self->queues = *queues; self->pushed = self->planned = self->watermark = 0; - int expected = 0; - if (!atomic_compare_exchange_strong(&metadata->active, &expected, 1)) - return DAMACY_INVAL; - expected = 0; - if (!atomic_compare_exchange_strong( - &metadata->reader->active, &expected, 1)) { - metadata->active = 0; - return DAMACY_INVAL; - } - self->bound = 1; self->reader = metadata_store_async_create( - (int)metadata->reader->concurrency, NULL, &metadata->reader->latency); + (int)metadata->reader.concurrency, NULL, &metadata->reader.latency); if (!self->reader) goto Fail; array_meta_async_fetcher_init(&self->array_fetcher, self->reader); @@ -313,7 +295,7 @@ damacy_chunk_planner_create(struct damacy_metadata* metadata, if (!self) return DAMACY_OOM; self->base.ops = &zarr_ops; - self->metadata = metadata; + self->metadata = *metadata; self->limits = *limits; *out = &self->base; return DAMACY_OK; diff --git a/tests/test_cpu_pipeline.c b/tests/test_cpu_pipeline.c index e5cd1c8a..5d170e50 100644 --- a/tests/test_cpu_pipeline.c +++ b/tests/test_cpu_pipeline.c @@ -223,6 +223,52 @@ test_owned_plan(void) return 0; } +static int +test_shared_metadata(void) +{ + char root[] = "/tmp/damacy_shared_metadata_XXXXXX"; + EXPECT(mkdtemp(root)); + char uri[256]; + snprintf(uri, sizeof(uri), "%s/array", root); + int64_t shape[] = { 5, 11 }, chunks[] = { 2, 4 }, shards[] = { 4, 8 }; + EXPECT(fixture_write_zarr(uri, shape, chunks, shards, 2, "uint16", 10) == 0); + struct components a = { 0 }, b = { 0 }; + EXPECT(create_components(&a) == 0); + EXPECT(create_components(&b) == 0); + damacy_planner_destroy(b.planner); + b.planner = NULL; + EXPECT(damacy_chunk_planner_create( + a.metadata, + &(struct damacy_plan_limits){ .max_chunks = 1024, + .max_chunk_bytes = 1 << 20, + .max_shards_per_sample = 4, + .max_plan_bytes = 1 << 20 }, + &b.planner) == DAMACY_OK); + EXPECT(start_pipeline(&a, DAMACY_F32, 2, 5, 2) == 0); + EXPECT(start_pipeline(&b, DAMACY_F32, 2, 5, 2) == 0); + damacy_metadata_reader_destroy(a.metadata_reader); + a.metadata_reader = NULL; + damacy_metadata_destroy(a.metadata); + a.metadata = NULL; + struct damacy* pipelines[] = { a.pipeline, b.pipeline }; + struct damacy_sample samples[] = { sample(uri, 1, 2, 2, 5), + sample(uri, 1, 2, 2, 5) }; + for (unsigned i = 0; i < 2; ++i) + EXPECT(damacy_push(pipelines[i], + (struct damacy_sample_slice){ samples, samples + 2 }) + .status == DAMACY_OK); + for (unsigned i = 0; i < 2; ++i) { + struct damacy_batch* batch = NULL; + EXPECT(damacy_pop(pipelines[i], &batch) == DAMACY_OK); + EXPECT(verify_crop(batch, 10, 16) == 0); + damacy_release(pipelines[i], batch); + } + destroy_components(&a); + destroy_components(&b); + fixture_rm_tree(root); + return 0; +} + struct pop_waiter { struct damacy* pipeline; @@ -348,6 +394,7 @@ main(void) { RUN(test_codecs_and_types); RUN(test_owned_plan); + RUN(test_shared_metadata); RUN(test_retained_outputs_and_shutdown); RUN(test_bfloat_rounding_and_fill); return 0; From 3353c1bffe335965137c7566372687071033a053 Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Thu, 24 Sep 2026 21:20:36 +0000 Subject: [PATCH 7/7] fix: batch API docs and release warning --- src/damacy.h | 17 ++++++++++++----- src/damacy_pop.c | 8 ++++++-- tests/test_cpu_pipeline.c | 32 ++++++++++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 7 deletions(-) diff --git a/src/damacy.h b/src/damacy.h index 077c1882..17bf9c15 100644 --- a/src/damacy.h +++ b/src/damacy.h @@ -194,8 +194,14 @@ extern "C" DAMACY_DEVICE_CUDA = 2, }; + // Add one batch reference. Each reference needs its own release. void damacy_batch_retain(struct damacy_batch* batch); + // Release one batch reference. Unlike damacy_release, this needs no + // pipeline, so it also works after the pipeline is destroyed. Thread-safe. void damacy_batch_release(struct damacy_batch* batch); + // Stop the pipeline without freeing it. Pending work is dropped, blocked + // damacy_pop callers wake with DAMACY_SHUTDOWN, and borrowed components may + // serve another pipeline. Retained batches stay valid. Safe to repeat. void damacy_shutdown(struct damacy* d); // Fill performance/resource knobs with explicit library defaults. Callers @@ -217,10 +223,10 @@ extern "C" // CUDA device index, or -1 for CPU execution. int damacy_get_device(const struct damacy* d); - // Tear down. Does NOT flush in-flight work; the io_queue is asked to - // shut down and pending CUDA streams are synchronized before buffers - // are released. Pending damacy_pop callers (from another thread) wake - // with DAMACY_SHUTDOWN. + // Shut down (see damacy_shutdown), then free the pipeline. Components + // passed to damacy_pipeline_create stay with the caller; components made by + // damacy_create are destroyed. Release retained batches afterwards with + // damacy_batch_release. void damacy_destroy(struct damacy* d); struct damacy_push_result @@ -250,7 +256,8 @@ extern "C" // Release one batch reference. The buffer is reusable after the last // consumer releases it. Thread-safe; may be called from - // a thread other than the one that called damacy_pop. + // a thread other than the one that called damacy_pop. A batch from another + // pipeline is still released, with a warning. void damacy_release(struct damacy* d, struct damacy_batch* b); // Release one reference after ordering CUDA output writes behind event. diff --git a/src/damacy_pop.c b/src/damacy_pop.c index e8f06dd5..286ab6cc 100644 --- a/src/damacy_pop.c +++ b/src/damacy_pop.c @@ -1,6 +1,7 @@ #include "damacy_internal.h" #include "damacy_stats.h" +#include "log/log.h" #include @@ -51,8 +52,11 @@ damacy_pop(struct damacy* self, struct damacy_batch** out) void damacy_release(struct damacy* self, struct damacy_batch* batch) { - if (batch && batch->owner == self) - damacy_batch_release(batch); + if (!batch) + return; + if (batch->owner != self) + log_warn("damacy_release: batch belongs to another pipeline"); + damacy_batch_release(batch); } enum damacy_status diff --git a/tests/test_cpu_pipeline.c b/tests/test_cpu_pipeline.c index 5d170e50..7791b014 100644 --- a/tests/test_cpu_pipeline.c +++ b/tests/test_cpu_pipeline.c @@ -269,6 +269,37 @@ test_shared_metadata(void) return 0; } +static int +test_release_from_other_pipeline(void) +{ + char root[] = "/tmp/damacy_release_XXXXXX"; + EXPECT(mkdtemp(root)); + char uri[256]; + snprintf(uri, sizeof(uri), "%s/array", root); + int64_t shape[] = { 5, 11 }, chunks[] = { 2, 4 }, shards[] = { 4, 8 }; + EXPECT(fixture_write_zarr(uri, shape, chunks, shards, 2, "uint16", 10) == 0); + struct components a = { 0 }, b = { 0 }; + EXPECT(create_components(&a) == 0); + EXPECT(create_components(&b) == 0); + EXPECT(start_pipeline(&a, DAMACY_F32, 2, 5, 2) == 0); + EXPECT(start_pipeline(&b, DAMACY_F32, 2, 5, 2) == 0); + struct damacy_sample samples[] = { sample(uri, 1, 2, 2, 5), + sample(uri, 1, 2, 2, 5) }; + EXPECT(damacy_push(a.pipeline, + (struct damacy_sample_slice){ samples, samples + 2 }) + .status == DAMACY_OK); + struct damacy_batch* batch = NULL; + EXPECT(damacy_pop(a.pipeline, &batch) == DAMACY_OK); + damacy_batch_retain(batch); + damacy_release(b.pipeline, batch); + EXPECT(atomic_load(&batch->references) == 1); + damacy_release(a.pipeline, batch); + destroy_components(&a); + destroy_components(&b); + fixture_rm_tree(root); + return 0; +} + struct pop_waiter { struct damacy* pipeline; @@ -395,6 +426,7 @@ main(void) RUN(test_codecs_and_types); RUN(test_owned_plan); RUN(test_shared_metadata); + RUN(test_release_from_other_pipeline); RUN(test_retained_outputs_and_shutdown); RUN(test_bfloat_rounding_and_fill); return 0;