diff --git a/README.md b/README.md index c9c8f675..41b42939 100644 --- a/README.md +++ b/README.md @@ -15,9 +15,9 @@ Both paths support raw bytes, zstd, and Blosc-zstd and return contiguous, DLPack-compatible tensors. CPU builds have no CUDA or nvCOMP dependency. 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. +[NGFF](https://ngff.openmicroscopy.org/) multiscale images. Queries select +rectangular regions or ordered index arrays along each dimension. Transformed +crops and automatic NGFF level selection are planned extensions. ## CPU quick start @@ -68,6 +68,23 @@ install dependencies with `brew install cmake ninja pkg-config zstd c-blosc`: pip install . --config-settings=cmake.define.DAMACY_CUDA=OFF ``` +## Indexed queries + +Use `IndexQuery` to mix index arrays and contiguous slices: + +```python +query = damacy.IndexQuery( + uri="/data/image.zarr/0", + selection=([7, 2, 7], slice(16, 80), [100, 3, 40, 3]), +) +assert query.shape == (3, 64, 4) +``` + +Choose `BatchSpec` with that sample shape and pass the query to `pipeline.push`. +Each axis is independent: all combinations of its selected indices appear in +the output. Order and duplicates are preserved. Both CPU and CUDA support the +same queries; see [query semantics](docs/pipeline.md#indexed-queries). + ## CUDA quick start ```python diff --git a/bench/main.c b/bench/main.c index afdcaeee..8b210d44 100644 --- a/bench/main.c +++ b/bench/main.c @@ -729,13 +729,15 @@ fill_random_sample(const struct scenario* sc, { uint32_t z = (uint32_t)rng_range(rng, at->n); s->uri = &at->uris[(size_t)z * BENCH_MAX_URI]; - s->aabb.rank = at->rank; + s->rank = at->rank; const int64_t* shape = &at->shapes[(size_t)z * DAMACY_MAX_RANK]; for (uint8_t d = 0; d < at->rank; ++d) { int64_t span = shape[d] - sc->sample_shape[d] + 1; int64_t beg = (int64_t)rng_range(rng, (uint64_t)span); - s->aabb.dims[d].beg = beg; - s->aabb.dims[d].end = beg + sc->sample_shape[d]; + s->axes[d] = (struct damacy_axis_selection){ + .kind = DAMACY_AXIS_INTERVAL, + .interval = { beg, beg + sc->sample_shape[d] } + }; } } diff --git a/dev/cpu-pipeline.md b/dev/cpu-pipeline.md index 6bd2784e..718bbcec 100644 --- a/dev/cpu-pipeline.md +++ b/dev/cpu-pipeline.md @@ -1,9 +1,9 @@ # 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. +CUDA executors. Both support rectangular and indexed queries. A CPU build +returns results in ordinary RAM and has no CUDA dependency. 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 @@ -16,7 +16,7 @@ existing callers. | 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`. | +| Planner | Validate requests, enumerate selected chunks, 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`. | @@ -79,10 +79,10 @@ adjacent or overlapping ranges within each shard and interleave merged reads across shards. Each unique chunk retains its offset within a merged read; its uses still control output assembly. CPU reads use exact encoded byte ranges, without CUDA's page alignment. It handles raw bytes, zstd, and C-Blosc zstd -with no, byte, or bit shuffle. It copies clipped chunk intersections and casts -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. +with no, byte, or bit shuffle. It copies clipped chunk intersections, gathers +indexed positions, and casts 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 @@ -115,20 +115,30 @@ 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. +## Indexed queries + +`IndexQuery` supplies ordered index vectors along selected dimensions and +contiguous slices along the others. Multiple indexed dimensions form a Cartesian +product. Order and repeated indices survive gathering; output shape is independent +of the bounding range. A stencil can generate an index vector without a separate +query type. + +`query/selection.c` copies and sorts each vector by source position while retaining +its output position. Preparation enumerates only touched grid cells, first for +shard metadata and then for inner chunks. A chunk has one use per sample even if +many selected positions, including duplicates, fall within it. `PLAN_GATHER` +regions own sorted index records along with the rest of the plan. Storage grows +with the sum of index-vector lengths and selected chunk uses. + +CPU assembly intersects each axis with a decoded chunk, then writes the Cartesian +product of selected positions to the specified output positions. Contiguous +trailing axes retain row copies. CUDA dispatch writes per-chunk selection spans +to a per-batch array that only indexed samples reference, so rectangular chunk +records carry no selection data. It uploads that array once per batch with +compact source/output index pairs. Assembly reuses the existing read, decode, +fill, shuffle, and cast paths. Index buffers hold the most indices a batch can +contain and are included in the GPU budget before wave sizing. A nonzero +`max_index_bytes` must cover them; zero rejects indexed samples at push. ## Next: spatial resampling and NGFF @@ -158,7 +168,7 @@ 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. +current chunk-at-a-time copy/gather 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. diff --git a/docs/api.md b/docs/api.md index 8ee83da1..6dd8a136 100644 --- a/docs/api.md +++ b/docs/api.md @@ -10,6 +10,8 @@ The full public surface of the `damacy` package. ::: damacy.Sample +::: damacy.IndexQuery + ::: damacy.Batch ## Components diff --git a/docs/budget.md b/docs/budget.md index 9e12b1e1..4d46153d 100644 --- a/docs/budget.md +++ b/docs/budget.md @@ -27,7 +27,8 @@ configuration — then round up generously: pool_reservation = 2 × samples_per_batch × prod(sample_shape) × dtype_bytes one_chunk_per_wave ≈ max_chunk_uncompressed_bytes × 2 # compressed + decoded buffers, one chunk both_waves_one_chunk ≈ 2 × one_chunk_per_wave # two GPU waves resident at once -budget_floor ≈ pool_reservation + both_waves_one_chunk + scratch_slack +index_reservation = 2 × (8 × samples_per_batch × sum(sample_shape) + 16 × 16384 × rank) # 0 if max_index_bytes = 0 +budget_floor ≈ pool_reservation + index_reservation + both_waves_one_chunk + scratch_slack ``` The `2 ×` in `pool_reservation` is output-side double-buffering. @@ -99,17 +100,23 @@ budget. ### Where the budget is spent -Four buckets share `max_gpu_memory_bytes`: +These allocations share `max_gpu_memory_bytes`: | Bucket | What it holds | Scales with | | ----------------------- | ---------------------------------------------------------- | -------------------------------------------------------- | | Output batch pool | The batches you `pop`, double-buffered | `samples_per_batch`, `sample_shape`, dtype | | Wave-resident buffers | Compressed + decoded chunk bytes for the two in-flight waves | budget headroom, in `max_chunk_uncompressed_bytes` steps | | Decoder scratch | nvcomp's working memory | Peak sub-stream count in the dataset | -| Per-wave metadata | Pointer/size arrays for the decoder | Peak sub-stream count | - -The first two are the large ones; the last two are small but -*depend on the data*. damacy cannot know the sub-stream count of +| Per-wave metadata | Pointer/size arrays for the decoder and assembly | Peak sub-stream and chunk counts | +| Indexed-query data | Source/output index pairs and per-chunk selection ranges for two batch slots | Output shape and rank; none when `max_index_bytes=0` | + +Index storage is reserved at construction from the output shape. A nonzero +`max_index_bytes` below `8 × samples_per_batch × sum(sample_shape)` raises +`BudgetExceeded` from `Pipeline(cfg)`. Each slot also reserves 16 bytes per +axis for up to 16384 chunks, to record which part of each chunk an indexed +query selects. Setting `max_index_bytes=0` drops both reservations; `push` +then rejects indexed queries. +Decoder scratch and fanout storage depend on the data. damacy cannot know the sub-stream count of a chunk until it inspects the chunk's header, so damacy picks per-wave geometry such that even after adaptive growth to the structural ceiling, the total fits inside the cap. Grows commit diff --git a/docs/pipeline.md b/docs/pipeline.md index a2a92a77..38c07a99 100644 --- a/docs/pipeline.md +++ b/docs/pipeline.md @@ -3,7 +3,7 @@ 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. +and indexed queries and prepared-plan contract. ## Construct a CPU pipeline @@ -96,10 +96,79 @@ 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. +## Indexed queries + +`IndexQuery(uri, selection=...)` accepts one index array or contiguous `slice` +per source axis. For a `(z, y, x)` array: + +```python +query = damacy.IndexQuery( + uri="/data/image.zarr/0", + selection=([7, 2, 7], slice(16, 80), [100, 3, 40, 3]), +) +output = damacy.BatchSpec(samples=1, shape=query.shape, dtype="f32") +assert query.shape == (3, 64, 4) + +with damacy.Pipeline( + planner=planner, executor=executor, output=output, + queues=damacy.QueueLimits(lookahead_samples=2), +) as pipeline: + pipeline.push([query]) + with pipeline.pop() as batch: + result = np.from_dlpack(batch) # CPU executor + del result +``` + +Index arrays select independently along each dimension: their Cartesian product +fills the tensor, matching MATLAB's per-dimension indexing or NumPy's `np.ix_`. +Caller order and repeated indices are preserved. A singleton index array keeps +its axis. A tuple such as `(7, 2)` is an index array in `IndexQuery`; use a slice +for a contiguous range. `Sample.aabb` retains its existing interval syntax. + +Indices must be nonnegative integers below `2**63 - 1`. Slices use half-open +bounds, require an explicit stop, and accept a step of one. A `range` object can +supply a strided or reversed index array. Empty selections are rejected. Source +bounds are validated when metadata arrives; out-of-bounds selections raise +`InvalidArgument` from `pop`. Each query's result shape must match `BatchSpec`. +Rectangular and indexed queries can share a batch when their result shapes match. + +Python copies index values into the immutable query at construction. Native +admission copies accepted queries, and prepared plans own their index data. +The planner enumerates only selected shards and chunks, including when indices +span large gaps. `max_chunks` counts each selected chunk once per sample; +repeated indices inside that chunk do not consume extra chunk entries. + +The C API uses `damacy_sample.rank` and one tagged `damacy_axis_selection` +per axis. C callers must rebuild against the updated headers and explicitly +set every active axis to `DAMACY_AXIS_INTERVAL` or `DAMACY_AXIS_INDICES`: + +```c +int64_t rows[] = {7, 2, 7}; +struct damacy_sample query = { + .uri = "/data/image.zarr/0", + .rank = 2, + .axes = { + {.kind = DAMACY_AXIS_INDICES, .indices = {.values = rows, .count = 3}}, + {.kind = DAMACY_AXIS_INTERVAL, .interval = {.beg = 4, .end = 8}}, + }, +}; +``` + +This requests a `(3, 4)` sample. The union's `kind` selects its active member; +there is no default kind. A zero or unknown tag, a null or empty index array, +or an interval with negative, empty, or reversed bounds returns `DAMACY_INVAL` +from `damacy_push`. Each axis's length must match the configured sample shape. +Only `axes[0..rank)` are active, in the Zarr array's stored axis order. Preparation +derives the bounding AABB from these selections. + +`damacy_push` copies the URI and index values for its consumed prefix. The +unconsumed suffix remains caller-owned and can be retried. + ## Limits and backpressure -Sizes are bytes, with positive explicit limits. Zero is invalid for these -capacities. Defaults come from the Python value objects. +Sizes are bytes, with positive explicit limits. `CudaLimits.max_index_bytes` +also accepts zero, which rejects indexed CUDA queries at push. Defaults come +from the Python value objects. | Setting | Scope | | --- | --- | @@ -110,13 +179,25 @@ capacities. Defaults come from the Python value objects. | `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. | +| `PlanLimits.max_plan_bytes` | Owned storage per prepared plan, including index arrays; also caps copied index data per queued query. | | `CpuLimits.max_memory_bytes` | CPU executor's buffers, codec workspace, active read plans, and temporary read-planning scratch. | | `CpuLimits.decode_workers` | Total decoding/assembly workers, including the calling scheduler thread. | | `CpuLimits.chunks_per_input_buffer` | Chunks each of the two encoded-input buffers holds; from `decode_workers` to 16384. | | `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. | +| `CudaLimits.max_index_bytes` | Device index storage per batch: eight bytes per index across all indexed axes and samples. Zero, or at least `8 * samples * sum(shape)`. Default 64 MiB. | + +CUDA reserves index storage for its two execution slots within +`max_gpu_memory_bytes`. A nonzero `max_index_bytes` must hold every index a +batch can contain, `8 * samples * sum(shape)` bytes, so an accepted query never +runs out of room. A smaller value raises `BudgetExceeded` from `Pipeline`. Each +slot allocates that amount, plus 16 bytes per axis for up to 16384 chunks to +record each chunk's selected range. An identical amount of pinned host staging +is allocated. With `max_index_bytes=0`, `push` raises `BudgetExceeded` for an +`IndexQuery`; the query is not consumed and the pipeline keeps running. +`Config.max_index_bytes` provides the same setting through the CUDA +convenience adapter. For injected pipelines, define `floor = queues.lookahead_samples + output.samples`. The metadata cache requires at least `floor` array entries and @@ -242,9 +323,8 @@ 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 +Rectangular and indexed queries copy or gather existing voxels. 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. diff --git a/python/damacy/__init__.py b/python/damacy/__init__.py index 8aec0c46..83539869 100644 --- a/python/damacy/__init__.py +++ b/python/damacy/__init__.py @@ -71,6 +71,7 @@ "DtypeMismatch", "FileMetadataReader", "FileReader", + "IndexQuery", "InvalidArgument", "LatencyModel", "MetadataCache", @@ -453,7 +454,82 @@ def __init__( ) def _to_native(self) -> dict[str, Any]: - return {"uri": self.uri, "aabb": list(self.aabb)} + return { + "uri": self.uri, + "axes": [("interval", bounds) for bounds in self.aabb], + } + + +@dataclass(init=False, frozen=True, slots=True) +class IndexQuery: + """Select independent index arrays or contiguous slices along each axis. + + Axes form a Cartesian product. Index order and duplicates are preserved; + singleton arrays keep their dimension. Indices are zero-based and must be + nonnegative. Slices require an explicit stop and a step of one. Empty + selections are rejected. Source bounds are checked when metadata arrives. + + >>> q = IndexQuery("cell.zarr", selection=([7, 2, 7], slice(4, 8))) + >>> q.shape + (3, 4) + >>> q.indices + ((7, 2, 7), None) + """ + + uri: str + aabb: tuple[tuple[int, int], ...] + indices: tuple[tuple[int, ...] | None, ...] + + def __init__(self, uri: str, selection: Iterable[slice | Iterable[int]]) -> None: + bounds: list[tuple[int, int]] = [] + indices: list[tuple[int, ...] | None] = [] + for axis, item in enumerate(selection): + if axis >= _native.MAX_RANK: + raise ValueError(f"selection rank must not exceed {_native.MAX_RANK}") + if isinstance(item, slice): + if item.step is not None and operator.index(item.step) != 1: + raise ValueError(f"selection axis {axis}: slice step must be one") + if item.stop is None: + raise ValueError(f"selection axis {axis}: slice stop is required") + start = 0 if item.start is None else operator.index(item.start) + stop = operator.index(item.stop) + if not 0 <= start < stop <= (1 << 63) - 1: + raise ValueError(f"selection axis {axis}: invalid slice bounds") + bounds.append((start, stop)) + indices.append(None) + else: + values = tuple(operator.index(value) for value in item) + if not values: + raise ValueError(f"selection axis {axis}: index array is empty") + if len(values) > (1 << 32) - 1: + raise ValueError(f"selection axis {axis}: too many indices") + if min(values) < 0 or max(values) >= (1 << 63) - 1: + raise ValueError( + f"selection axis {axis}: indices must be nonnegative int64 values" + ) + bounds.append((min(values), max(values) + 1)) + indices.append(values) + if not bounds: + raise ValueError("selection must have at least one axis") + object.__setattr__(self, "uri", uri) + object.__setattr__(self, "aabb", tuple(bounds)) + object.__setattr__(self, "indices", tuple(indices)) + + @property + def shape(self) -> tuple[int, ...]: + return tuple( + len(indices) if indices is not None else end - begin + for (begin, end), indices in zip(self.aabb, self.indices, strict=True) + ) + + def _to_native(self) -> dict[str, Any]: + return { + "uri": self.uri, + "axes": [ + ("indices", indices) if indices is not None else ("interval", bounds) + for bounds, indices in zip(self.aabb, self.indices, strict=True) + ], + } @dataclass(init=False, frozen=True, slots=True) @@ -519,6 +595,10 @@ class Config: intersecting more shards is rejected with :class:`InvalidArgument`. max_chunk_uncompressed_bytes: Largest uncompressed chunk size the pipeline accepts; 0 selects the C default (512 KB). + max_index_bytes: GPU index storage per batch, eight bytes per index. + Must be zero, which rejects indexed queries at push, or at least + ``8 * samples_per_batch * sum(sample_shape)``. Included in the + total GPU budget. max_read_op_bytes: Cap on the size of a single coalesced read issued to storage. 0 selects the C default. Tune against your storage tier: small values keep the queue @@ -562,6 +642,7 @@ class Config: max_chunk_uncompressed_bytes: int max_read_op_bytes: int max_gpu_memory_bytes: int + max_index_bytes: int host_buffer_waves: int max_chunks_per_wave: int max_substreams_per_chunk: int @@ -589,6 +670,7 @@ def __init__( lookahead_samples: int | None = None, max_chunk_uncompressed_bytes: int = _native.DEFAULT_CHUNK_UNCOMPRESSED_BYTES, max_read_op_bytes: int = _native.DEFAULT_READ_OP_MAX_BYTES, + max_index_bytes: int = _native.DEFAULT_MAX_INDEX_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, @@ -639,6 +721,7 @@ def __init__( "max_read_op_bytes must be in " f"[1, {_native.MAX_READ_OP_BYTES}] (got {max_read_op_bytes})" ) + max_index_bytes = _index_byte_limit(max_index_bytes) if max_gpu_memory_bytes < 1: raise ValueError( f"max_gpu_memory_bytes must be >= 1 (got {max_gpu_memory_bytes})" @@ -693,6 +776,7 @@ def __init__( set_(self, "lookahead_samples", lookahead_samples) set_(self, "max_chunk_uncompressed_bytes", max_chunk_uncompressed_bytes) set_(self, "max_read_op_bytes", max_read_op_bytes) + set_(self, "max_index_bytes", max_index_bytes) set_(self, "max_gpu_memory_bytes", max_gpu_memory_bytes) set_(self, "host_buffer_waves", host_buffer_waves) set_(self, "max_chunks_per_wave", max_chunks_per_wave) @@ -736,6 +820,13 @@ def default( ) +def _index_byte_limit(value: int) -> int: + value = operator.index(value) + if not 0 <= value <= 8 * ((1 << 32) - 1): + raise ValueError("max_index_bytes must be between 0 and 8 * (2**32 - 1)") + return value + + def _positive_int(value: int, name: str, maximum: int = (1 << 32) - 1) -> int: value = operator.index(value) if not 1 <= value <= maximum: @@ -843,9 +934,18 @@ def __post_init__(self) -> None: @dataclass(frozen=True, slots=True) class CudaLimits: - """GPU buffer limits, wave geometry, and codec-layout cache capacity.""" + """GPU buffer limits, wave geometry, and codec-layout cache capacity. + + ``max_index_bytes`` is index storage per batch, eight bytes per index. It + must be zero or at least ``8 * samples * sum(shape)`` of the output, so a + batch never runs out; a smaller value raises :class:`BudgetExceeded` from + :class:`Pipeline`. With zero, :meth:`Pipeline.push` rejects an + :class:`IndexQuery` with :class:`BudgetExceeded`. Storage counts against + the total GPU budget; only what the output shape needs is allocated. + """ max_gpu_memory_bytes: int + max_index_bytes: int = _native.DEFAULT_MAX_INDEX_BYTES chunk_layout_entries: int = 256 max_chunk_bytes: int = _native.DEFAULT_CHUNK_UNCOMPRESSED_BYTES max_read_bytes: int = _native.DEFAULT_READ_OP_MAX_BYTES @@ -854,6 +954,7 @@ class CudaLimits: max_substreams_per_chunk: int = _native.DEFAULT_MAX_SUBSTREAMS_PER_CHUNK def __post_init__(self) -> None: + _index_byte_limit(self.max_index_bytes) _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") @@ -1005,6 +1106,7 @@ def __init__( int(NumaStrategy.coerce(numa_strategy)), numa_node, _gds_to_native(enable_gds), + limits.max_index_bytes, ) @@ -1455,6 +1557,7 @@ def __init__( 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_index_bytes=config.max_index_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, @@ -1506,8 +1609,8 @@ def __init__( # _pending_buf is the head iterator's already-pulled-but-not-yet- # pushed samples; held flat to avoid wrapping `it` in successive # itertools.chain() layers under sustained backpressure. - self._pending: deque[Iterator[Sample]] = deque() - self._pending_buf: list[Sample] = [] + self._pending: deque[Iterator[Sample | IndexQuery]] = deque() + self._pending_buf: list[Sample | IndexQuery] = [] # damacy_pop has no timed variant; on timeout the worker stays # parked inside it and the next pop() adopts the same thread. self._pop_lock = threading.Lock() @@ -1588,15 +1691,17 @@ def __exit__( # ---- pipeline ---------------------------------------------------- - def push(self, samples: Iterable[Sample]) -> None: + def push(self, samples: Iterable[Sample | IndexQuery]) -> None: """Queue samples for processing. Accepts any iterable (list, generator, infinite generator, …); large or unbounded sources are pulled lazily as :meth:`pop` frees space. 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`, + discards the offending iterator. So does an :class:`IndexQuery` + on a CUDA executor with ``max_index_bytes=0``, which raises + :class:`BudgetExceeded`; the pipeline keeps running. Errors that + depend on store contents — :class:`NotFound`, :class:`DtypeMismatch`, per-array :class:`RankMismatch`, decode failures — surface at :meth:`pop` instead, since the pipeline fetches metadata asynchronously after push returns. Once any such error fires, diff --git a/python/damacy/_api.c b/python/damacy/_api.c index 621f6706..a9793c25 100644 --- a/python/damacy/_api.c +++ b/python/damacy/_api.c @@ -739,6 +739,7 @@ Pipeline_init(PipelineObj* self, PyObject* args, PyObject* kw) "metadata_latency_lognormal_sigma_ln_ns", "metadata_latency_cap_ns", "metadata_latency_seed", + "max_index_bytes", NULL }; struct damacy_tuning td = damacy_tuning_defaults(); unsigned int samples_per_batch = 0; @@ -746,6 +747,7 @@ Pipeline_init(PipelineObj* self, PyObject* args, PyObject* kw) PyObject* dtype_obj = NULL; unsigned int max_chunk_uncompressed = td.max_chunk_uncompressed_bytes; unsigned long long max_gpu_bytes = 0; + unsigned long long max_index_bytes = td.max_index_bytes; unsigned int n_io = td.n_io_threads; unsigned int metadata_io_concurrency = td.metadata_io_concurrency; unsigned int n_array_meta = DAMACY_DEFAULT_ARRAY_META_CACHE; @@ -770,7 +772,7 @@ Pipeline_init(PipelineObj* self, PyObject* args, PyObject* kw) unsigned long long metadata_latency_seed = 0; if (!PyArg_ParseTupleAndKeywords(args, kw, - "IIOIKIIIIIIO|IIIKiiiipKddKK", + "IIOIKIIIIIIO|IIIKiiiipKddKKK", kws, &samples_per_batch, &lookahead, @@ -797,7 +799,8 @@ Pipeline_init(PipelineObj* self, PyObject* args, PyObject* kw) &metadata_latency_lognormal_mu_ln_ns, &metadata_latency_lognormal_sigma_ln_ns, &metadata_latency_cap_ns, - &metadata_latency_seed)) + &metadata_latency_seed, + &max_index_bytes)) return -1; if (enable_gds != DAMACY_GDS_AUTO && enable_gds != DAMACY_GDS_ON && @@ -833,6 +836,7 @@ Pipeline_init(PipelineObj* self, PyObject* args, PyObject* kw) .device = device, .tuning = { .max_gpu_memory_bytes = (uint64_t)max_gpu_bytes, + .max_index_bytes = (uint64_t)max_index_bytes, .max_chunk_uncompressed_bytes = max_chunk_uncompressed, .max_read_op_bytes = (uint64_t)max_read_op_bytes, .host_buffer_waves = (uint8_t)host_buffer_waves, @@ -927,42 +931,103 @@ parse_sample(PyObject* obj, struct damacy_sample* out) PyErr_SetString(PyExc_TypeError, "sample must be a dict"); return -1; } - PyObject* uri = PyDict_GetItemString(obj, "uri"); // borrowed - PyObject* aabb = PyDict_GetItemString(obj, "aabb"); // borrowed - if (!uri || !aabb) { - PyErr_SetString(PyExc_KeyError, "sample requires 'uri' and 'aabb'"); + PyObject* uri = PyDict_GetItemString(obj, "uri"); + PyObject* axes = PyDict_GetItemString(obj, "axes"); + if (!uri || !axes) { + PyErr_SetString(PyExc_KeyError, "sample requires 'uri' and 'axes'"); + return -1; + } + if (PyDict_Size(obj) != 2) { + PyErr_SetString(PyExc_ValueError, "sample accepts only 'uri' and 'axes'"); return -1; } const char* uri_s = PyUnicode_AsUTF8(uri); if (!uri_s) return -1; - if (!PyList_Check(aabb) && !PyTuple_Check(aabb)) { - PyErr_SetString(PyExc_TypeError, - "aabb must be a list or tuple of (beg,end)"); + if (!PyList_Check(axes) && !PyTuple_Check(axes)) { + PyErr_SetString(PyExc_TypeError, "axes must be a list or tuple"); return -1; } - Py_ssize_t n = PySequence_Fast_GET_SIZE(aabb); + Py_ssize_t n = PySequence_Fast_GET_SIZE(axes); if (n < 1 || n > DAMACY_MAX_RANK) { - PyErr_Format( - PyExc_ValueError, "aabb rank out of range: %zd", (Py_ssize_t)n); + PyErr_Format(PyExc_ValueError, "sample rank out of range: %zd", n); return -1; } out->uri = uri_s; - out->aabb.rank = (uint8_t)n; - for (Py_ssize_t i = 0; i < n; ++i) { - PyObject* item = PySequence_Fast_GET_ITEM(aabb, i); - long long beg, end; - if (!PyArg_ParseTuple(item, "LL", &beg, &end)) { - PyErr_Format(PyExc_ValueError, "aabb[%zd] must be a (beg,end) pair", i); + out->rank = (uint8_t)n; + for (Py_ssize_t d = 0; d < n; ++d) { + PyObject* item = PySequence_Fast_GET_ITEM(axes, d); + if ((!PyList_Check(item) && !PyTuple_Check(item)) || + PySequence_Fast_GET_SIZE(item) != 2) { + PyErr_Format( + PyExc_ValueError, "axes[%zd] must be a (kind, values) pair", d); + return -1; + } + PyObject* kind = PySequence_Fast_GET_ITEM(item, 0); + PyObject* values = PySequence_Fast_GET_ITEM(item, 1); + if (!PyUnicode_Check(kind)) { + PyErr_Format(PyExc_TypeError, "axes[%zd] kind must be a string", d); + return -1; + } + struct damacy_axis_selection* axis = &out->axes[d]; + if (PyUnicode_CompareWithASCIIString(kind, "interval") == 0) { + long long beg, end; + if (!PyArg_ParseTuple(values, "LL", &beg, &end)) { + PyErr_Format( + PyExc_ValueError, "axes[%zd] interval must be a (beg, end) pair", d); + return -1; + } + *axis = (struct damacy_axis_selection){ + .kind = DAMACY_AXIS_INTERVAL, .interval = { (int64_t)beg, (int64_t)end } + }; + } else if (PyUnicode_CompareWithASCIIString(kind, "indices") == 0) { + PyObject* sequence = + PySequence_Fast(values, "index array must be iterable"); + if (!sequence) + return -1; + Py_ssize_t count = PySequence_Fast_GET_SIZE(sequence); + if (!count || (uint64_t)count > UINT32_MAX || + (size_t)count > SIZE_MAX / sizeof(int64_t)) { + Py_DECREF(sequence); + PyErr_SetString(PyExc_ValueError, "index array size is out of range"); + return -1; + } + int64_t* copied = PyMem_Malloc((size_t)count * sizeof(*copied)); + if (!copied) { + Py_DECREF(sequence); + PyErr_NoMemory(); + return -1; + } + *axis = (struct damacy_axis_selection){ + .kind = DAMACY_AXIS_INDICES, .indices = { copied, (uint32_t)count } + }; + for (Py_ssize_t i = 0; i < count; ++i) { + copied[i] = PyLong_AsLongLong(PySequence_Fast_GET_ITEM(sequence, i)); + if (PyErr_Occurred()) { + Py_DECREF(sequence); + return -1; + } + } + Py_DECREF(sequence); + } else { + PyErr_Format(PyExc_ValueError, "axes[%zd] has unknown kind %R", d, kind); return -1; } - out->aabb.dims[i].beg = (int64_t)beg; - out->aabb.dims[i].end = (int64_t)end; } return 0; } +static void +free_samples(struct damacy_sample* samples, Py_ssize_t count) +{ + for (Py_ssize_t i = 0; i < count; ++i) + for (uint8_t d = 0; d < samples[i].rank; ++d) + if (samples[i].axes[d].kind == DAMACY_AXIS_INDICES) + PyMem_Free((void*)samples[i].axes[d].indices.values); + PyMem_Free(samples); +} + static PyObject* Pipeline_push(PipelineObj* self, PyObject* arg) { @@ -983,7 +1048,7 @@ Pipeline_push(PipelineObj* self, PyObject* arg) for (Py_ssize_t i = 0; i < n; ++i) { PyObject* item = PySequence_Fast_GET_ITEM(arg, i); if (parse_sample(item, &buf[i]) != 0) { - PyMem_Free(buf); + free_samples(buf, n); return NULL; } } @@ -993,7 +1058,7 @@ Pipeline_push(PipelineObj* self, PyObject* arg) WITH_GIL_RELEASED(r = damacy_push(self->handle, slice)); Py_ssize_t consumed = (Py_ssize_t)(r.unconsumed.beg - buf); - PyMem_Free(buf); + free_samples(buf, n); // OK / AGAIN are not errors at this layer — return the consumed count // and the integer status so the caller can detect back-pressure. Any @@ -1302,6 +1367,7 @@ api_register_types(PyObject* m) } limits[] = { { "DEFAULT_CHUNK_UNCOMPRESSED_BYTES", d.max_chunk_uncompressed_bytes }, { "DEFAULT_READ_OP_MAX_BYTES", d.max_read_op_bytes }, + { "DEFAULT_MAX_INDEX_BYTES", d.max_index_bytes }, { "DEFAULT_HOST_BUFFER_WAVES", d.host_buffer_waves }, { "DEFAULT_MAX_CHUNKS_PER_WAVE", d.max_chunks_per_wave }, { "DEFAULT_MAX_SUBSTREAMS_PER_CHUNK", d.max_substreams_per_chunk }, @@ -1311,6 +1377,7 @@ api_register_types(PyObject* m) { "DEFAULT_SHARD_INDEX_CACHE", d.n_shard_index_cache }, { "DEFAULT_CHUNK_LAYOUT_CACHE", d.n_chunk_layout_cache }, { "DEFAULT_MAX_SHARDS_PER_SAMPLE", d.max_shards_per_sample }, + { "MAX_RANK", DAMACY_MAX_RANK }, { "MAX_CHUNK_BYTES", DAMACY_MAX_CHUNK_BYTES }, { "MAX_READ_OP_BYTES", UINT32_MAX }, { "N_WAVES", DAMACY_N_WAVES }, diff --git a/python/damacy/_components.c b/python/damacy/_components.c index d357a30b..2753a2c2 100644 --- a/python/damacy/_components.c +++ b/python/damacy/_components.c @@ -211,12 +211,12 @@ create_cuda_executor(PyObject* self, PyObject* args) { (void)self; PyObject* reader_object; - unsigned long long memory, read_bytes; + unsigned long long memory, read_bytes, index_bytes; unsigned int buffers; int numa, gds; struct damacy_cuda_config config = { 0 }; if (!PyArg_ParseTuple(args, - "OiKIKIIIIiii", + "OiKIKIIIIiiiK", &reader_object, &config.device, &memory, @@ -228,13 +228,15 @@ create_cuda_executor(PyObject* self, PyObject* args) &config.chunk_layout_entries, &numa, &config.numa_node, - &gds)) + &gds, + &index_bytes)) 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_index_bytes = index_bytes; config.max_read_bytes = read_bytes; config.host_buffer_waves = (uint8_t)buffers; config.numa_strategy = (enum damacy_numa_strategy)numa; diff --git a/python/damacy/_native.pyi b/python/damacy/_native.pyi index b7a5ad24..03f432b7 100644 --- a/python/damacy/_native.pyi +++ b/python/damacy/_native.pyi @@ -7,6 +7,8 @@ from typing import Any, Final __version__: Final[str] CUDA_ENABLED: Final[int] +MAX_RANK: Final[int] +DEFAULT_MAX_INDEX_BYTES: Final[int] # ---- log-level constants (mirror damacy_log.h) -------------------------- @@ -154,13 +156,15 @@ class Pipeline: metadata_latency_lognormal_sigma_ln_ns: float = 0.0, metadata_latency_cap_ns: int = 0, metadata_latency_seed: int = 0, + max_index_bytes: int = 64 << 20, ) -> None: ... @property def device(self) -> int: ... def push( self, samples: list[dict[str, Any]] | tuple[dict[str, Any], ...] ) -> dict[str, Any]: - """Push a sequence of {uri, aabb} dicts. + """Push {uri, axes} dicts. Each axis is ("interval", (beg, end)) + or ("indices", values). Unknown kinds and extra fields are rejected. Returns {consumed: int, status: int}. Raises DamacyError for anything other than OK / AGAIN.""" @@ -235,6 +239,7 @@ def create_cuda_executor( numa_strategy: int, numa_node: int, gds: int, + max_index_bytes: int, /, ) -> object: ... def compose_pipeline( diff --git a/python/tests/test_index_queries.py b/python/tests/test_index_queries.py new file mode 100644 index 00000000..2359135a --- /dev/null +++ b/python/tests/test_index_queries.py @@ -0,0 +1,458 @@ +from __future__ import annotations + +import ctypes +import dataclasses +import json +import subprocess +from pathlib import Path + +import damacy +import numpy as np +import pytest + + +def planner(**overrides): + 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, + } + | overrides + ) + ), + ) + + +def cpu_executor(): + return damacy.CpuExecutor( + reader=damacy.FileReader(workers=2, max_inflight_reads=1), + limits=damacy.CpuLimits( + max_memory_bytes=8 << 20, + decode_workers=2, + max_encoded_chunk_bytes=1024, + max_decoded_chunk_bytes=1024, + ), + ) + + +def pipeline(*, shape, samples=1, **overrides): + return damacy.Pipeline( + **( + { + "planner": planner(), + "executor": cpu_executor(), + "output": damacy.BatchSpec(samples, shape), + "queues": damacy.QueueLimits(lookahead_samples=4), + "pop_timeout_s": 5.0, + } + | overrides + ) + ) + + +@pytest.fixture(params=["cpu", "cuda"]) +def indexed_executor(request): + if request.param == "cpu": + return cpu_executor() + request.getfixturevalue("cuda_ctx") + return damacy.CudaExecutor( + reader=damacy.FileReader(workers=2, max_inflight_reads=1), + device=0, + limits=damacy.CudaLimits( + max_gpu_memory_bytes=256 << 20, + max_chunk_bytes=1024, + max_chunks_per_wave=2, + max_substreams_per_chunk=8, + ), + ) + + +def read_batch(batch): + dtype = np.float32 if batch.info.dtype == damacy.Dtype.F32 else np.uint16 + result = np.empty(batch.info.shape, dtype=dtype) + if batch.info.device_type == damacy.DeviceType.CPU: + ctypes.memmove(result.ctypes.data, batch.info.data, result.nbytes) + else: + 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 + assert copy(result.ctypes.data, batch.info.data, result.nbytes) == 0 + if dtype == np.uint16: + return (result.astype(np.uint32) << 16).view(np.float32) + return result + + +def round_bfloat(values): + bits = values.astype(np.float32).view(np.uint32) + return ((bits + 0x7FFF + ((bits >> 16) & 1)) & 0xFFFF0000).view(np.float32) + + +def raw_chunks(path, values, chunks, *, fill=0, missing=()): + 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": chunks}, + }, + "chunk_key_encoding": { + "name": "default", + "configuration": {"separator": "/"}, + }, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + } + ) + ) + grid = tuple((n + c - 1) // c for n, c in zip(values.shape, chunks, strict=True)) + for coordinate in np.ndindex(grid): + if coordinate in missing: + continue + slices = tuple( + slice(i * c, min((i + 1) * c, n)) + for i, c, n in zip(coordinate, chunks, values.shape, strict=True) + ) + block = np.full(chunks, fill, dtype=values.dtype) + selected = values[slices] + block[tuple(slice(0, n) for n in selected.shape)] = selected + file = path.joinpath("c", *map(str, coordinate)) + file.parent.mkdir(parents=True, exist_ok=True) + file.write_bytes(block.tobytes()) + return str(path) + + +def test_native_axes_reject_malformed_queries(tmp_path): + values = np.arange(16, dtype=np.float32).reshape(4, 4) + uri = raw_chunks(tmp_path / "array", values, (2, 2)) + invalid = [ + (("unknown", (1, 3)), ValueError), + ((None, (1, 3)), TypeError), + ((), ValueError), + (("interval", (1, 3), "indices"), ValueError), + (("interval", (1, 2, 3)), ValueError), + (("interval", None), ValueError), + (("indices", ()), ValueError), + (("indices", None), TypeError), + (("indices", (1, 1.5)), TypeError), + ] + with pipeline(shape=(2, 2)) as p: + for axis, error in invalid: + with pytest.raises(error): + p._native.push([{"uri": uri, "axes": [("indices", (3, 1)), axis]}]) + with pytest.raises(ValueError, match="only 'uri' and 'axes'"): + p._native.push( + [{"uri": uri, "axes": [("interval", (0, 2))] * 2, "aabb": [(0, 2)] * 2}] + ) + with pytest.raises(KeyError, match="requires 'uri' and 'axes'"): + p._native.push( + [{"uri": uri, "aabb": [(0, 2)] * 2, "indices": [(1, 0)] * 2}] + ) + p.push([damacy.IndexQuery(uri, ([3, 1], slice(1, 3)))]) + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch)[0], values[[3, 1], 1:3]) + + +def test_query_copies_and_normalises_inputs(): + rows = np.array([7, 2, 7], dtype=np.int32) + q = damacy.IndexQuery("a", (rows, slice(None, 4))) + rows[:] = 0 + assert q.indices == ((7, 2, 7), None) + assert q.shape == (3, 4) + assert q == damacy.IndexQuery("a", ((7, 2, 7), slice(0, 4, 1))) + assert hash(q) == hash(damacy.IndexQuery("a", ([7, 2, 7], slice(0, 4)))) + + +@pytest.mark.parametrize( + "selection", + [ + (), + ([],), + ([-1],), + ([2**63 - 1],), + (slice(0, 0),), + (slice(2, 1),), + (slice(None),), + (slice(0, 4, 2),), + ], +) +def test_invalid_query_values(selection): + with pytest.raises(ValueError): + damacy.IndexQuery("a", selection) + + +@pytest.mark.parametrize("selection", [([1.5],), (["1"],), (1,), (slice(0.5, 2),)]) +def test_non_integer_indices(selection): + with pytest.raises(TypeError): + damacy.IndexQuery("a", selection) + + +@pytest.mark.parametrize( + "dtype", ["uint8", "uint16", "int16", "uint32", "int32", "float16", "float32"] +) +@pytest.mark.parametrize("output_dtype", ["f32", "bf16"]) +def test_index_order_duplicates_and_edge_chunks( + indexed_executor, tmp_path, dtype, output_dtype +): + values = (np.arange(55).reshape(5, 11) * 13).astype(dtype) + uri = raw_chunks(tmp_path / "array", values, (2, 4)) + rows, cols = [4, 0, 4, 1], [10, 3, 3, 0, 7] + expected = values[np.ix_(rows, cols)].astype(np.float32) + if output_dtype == "bf16": + expected = round_bfloat(expected) + with pipeline( + shape=(4, 5), + executor=indexed_executor, + planner=planner(max_shards_per_sample=9), + output=damacy.BatchSpec(1, (4, 5), output_dtype), + ) as p: + p.push([damacy.IndexQuery(uri, (rows, cols))]) + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch), expected[None]) + + +@pytest.mark.parametrize( + "selection", + [ + ([4, 0, 4], slice(2, 7)), + (slice(1, 4), [10, 2, 10, 0, 5]), + ([4], [10]), + (slice(1, 4), slice(2, 7)), + ], +) +def test_mixed_axes(indexed_executor, tmp_path, selection): + values = np.arange(55, dtype=np.float32).reshape(5, 11) + uri = raw_chunks(tmp_path / "array", values, (2, 4)) + query = damacy.IndexQuery(uri, selection) + axes = [ + np.arange(n)[axis] if isinstance(axis, slice) else axis + for n, axis in zip(values.shape, selection, strict=True) + ] + expected = values[np.ix_(*axes)] + with pipeline( + shape=query.shape, + executor=indexed_executor, + planner=planner(max_shards_per_sample=9), + ) as p: + p.push([query]) + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch), expected[None]) + + +def test_sparse_queries_only_plan_touched_chunks(indexed_executor, tmp_path): + values = np.arange(1024, dtype=np.float32).reshape(32, 32) + uri = raw_chunks(tmp_path / "array", values, (8, 8)) + Path(uri, "c", "1", "1").write_bytes(b"unselected corrupt chunk") + indices = [31, 0, 31] + with pipeline( + shape=(3, 3), executor=indexed_executor, planner=planner(max_chunks=4) + ) as p: + p.push([damacy.IndexQuery(uri, (indices, indices))]) + with p.pop() as batch: + np.testing.assert_array_equal( + read_batch(batch), values[np.ix_(indices, indices)][None] + ) + assert p.stats().chunks_planned == 4 + assert p.stats().assemble.output_bytes == 36 + + +def test_mixed_queries_batches_and_repeated_indices(indexed_executor, tmp_path): + values = np.arange(16, dtype=np.float32).reshape(4, 4) + uri = raw_chunks(tmp_path / "array", values, (2, 2)) + query = damacy.IndexQuery(uri, ([3, 0, 3, 0], [3, 3, 0, 0])) + rect = damacy.Sample(uri, [(0, 4), (0, 4)]) + expected = np.stack([values[np.ix_([3, 0, 3, 0], [3, 3, 0, 0])], values]) + with pipeline(shape=(4, 4), samples=2, executor=indexed_executor) as p: + p.push(q for _ in range(5) for q in (query, rect)) + for _ in range(5): + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch), expected) + + +def test_repeats_can_exceed_source_extent(indexed_executor, tmp_path): + values = np.array([[3.25, -7.5]], dtype=np.float32) + uri = raw_chunks(tmp_path / "array", values, (1, 2)) + query = damacy.IndexQuery(uri, ([0] * 17, [1, 0] * 31)) + with pipeline(shape=query.shape, executor=indexed_executor) as p: + p.push([query]) + with p.pop() as batch: + np.testing.assert_array_equal( + read_batch(batch), np.tile(values[:, ::-1], (17, 31))[None] + ) + + +def test_missing_chunks(indexed_executor, tmp_path): + values = np.arange(16, dtype=np.float32).reshape(4, 4) + uri = raw_chunks( + tmp_path / "array", values, (2, 2), fill=-7, missing=((0, 1), (1, 0)) + ) + values[:2, 2:] = -7 + values[2:, :2] = -7 + indices = [3, 0, 3] + with pipeline(shape=(3, 3), executor=indexed_executor) as p: + p.push([damacy.IndexQuery(uri, (indices, indices))]) + with p.pop() as batch: + np.testing.assert_array_equal( + read_batch(batch), values[np.ix_(indices, indices)][None] + ) + + +@pytest.mark.parametrize( + "codec,shuffle", + [ + ("zstd", "noshuffle"), + ("blosc-zstd", "noshuffle"), + ("blosc-zstd", "shuffle"), + ("blosc-zstd", "bitshuffle"), + ], +) +def test_compressed_shards( + indexed_executor, tmp_path, write_zarr_script, codec, shuffle +): + uri = tmp_path / "array" + subprocess.run( + [ + "uv", + "run", + "--script", + str(write_zarr_script), + "--out", + str(uri), + "--shape", + "5,11", + "--inner", + "2,4", + "--shard", + "4,8", + "--dtype", + "float32", + "--codec", + codec, + "--shuffle", + shuffle, + ], + check=True, + capture_output=True, + ) + rows, cols = [4, 0, 3, 4], [10, 0, 3, 10] + expected = np.arange(55, dtype=np.float32).reshape(5, 11)[np.ix_(rows, cols)] + with pipeline(shape=(4, 4), executor=indexed_executor) as p: + p.push([damacy.IndexQuery(str(uri), (rows, cols))]) + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch), expected[None]) + + +@pytest.mark.parametrize("rank", [1, 3, 8, 9, 31]) +def test_indexed_ranks(indexed_executor, tmp_path, rank): + shape = (1,) * (rank - 1) + (7,) + values = np.arange(7, dtype=np.float32).reshape(shape) + uri = raw_chunks(tmp_path / "array", values, (1,) * (rank - 1) + (3,)) + selection = ([0],) * (rank - 1) + ([6, 0, 6, 2],) + query = damacy.IndexQuery(uri, selection) + with pipeline(shape=query.shape, executor=indexed_executor) as p: + p.push([query]) + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch).ravel(), [6, 0, 6, 2]) + + +def test_index_bounds_and_result_shape(tmp_path): + uri = raw_chunks( + tmp_path / "array", np.arange(4, dtype=np.float32).reshape(2, 2), (2, 2) + ) + with pipeline(shape=(1, 1)) as p, pytest.raises(damacy.InvalidArgument): + p.push([damacy.IndexQuery(uri, ([0, 1], [0]))]) + with pipeline(shape=(1, 1)) as p: + p.push([damacy.IndexQuery(uri, ([2], [0]))]) + with pytest.raises(damacy.InvalidArgument): + p.pop() + with ( + pipeline(shape=(1, 1), planner=planner(max_plan_bytes=16)) as p, + pytest.raises(damacy.BudgetExceeded), + ): + p.push([damacy.IndexQuery(uri, ([0], [0]))]) + + +def test_large_sparse_coordinates(indexed_executor, tmp_path): + uri = raw_chunks(tmp_path / "array", np.array([3, 7], dtype=np.float32), (1,)) + far = 1 << 40 + metadata = Path(uri, "zarr.json") + config = json.loads(metadata.read_text()) + config["shape"] = [far + 1] + metadata.write_text(json.dumps(config)) + Path(uri, "c", "1").rename(Path(uri, "c", str(far))) + with pipeline( + shape=(3,), + executor=indexed_executor, + planner=planner(max_chunks=2, max_shards_per_sample=2), + ) as p: + p.push([damacy.IndexQuery(uri, ([far, 0, far],))]) + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch), [[7, 3, 7]]) + + +def test_random_cartesian_selections(indexed_executor, tmp_path): + values = np.arange(5 * 6 * 7, dtype=np.float32).reshape(5, 6, 7) + uri = raw_chunks(tmp_path / "array", values, (3, 3, 4)) + rng = np.random.default_rng(614) + selections = [ + tuple(rng.integers(0, n, size=4) for n in values.shape) for _ in range(12) + ] + with pipeline( + shape=(4, 4, 4), + samples=2, + executor=indexed_executor, + planner=planner(max_shards_per_sample=8), + ) as p: + p.push(damacy.IndexQuery(uri, axes) for axes in selections) + for i in range(0, len(selections), 2): + expected = np.stack( + [values[np.ix_(*axes)] for axes in selections[i : i + 2]] + ) + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch), expected) + + +def test_cuda_index_budget(indexed_executor, tmp_path): + if not isinstance(indexed_executor, damacy.CudaExecutor): + pytest.skip("CUDA index storage limit") + uri = raw_chunks( + tmp_path / "array", np.arange(4, dtype=np.float32).reshape(2, 2), (2, 2) + ) + rectangle = damacy.Sample(uri, [(0, 2), (0, 2)]) + query = damacy.IndexQuery(uri, ([1, 0], [1, 0])) + + def executor(limit): + return damacy.CudaExecutor( + reader=indexed_executor.reader, + limits=dataclasses.replace(indexed_executor.limits, max_index_bytes=limit), + device=0, + ) + + with pytest.raises(damacy.BudgetExceeded): + pipeline(shape=(2, 2), executor=executor(31)) + with pipeline(shape=(2, 2), executor=executor(0)) as p: + with pytest.raises(damacy.BudgetExceeded): + p.push([rectangle, query]) + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch), [[[0, 1], [2, 3]]]) + p.push([rectangle]) + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch), [[[0, 1], [2, 3]]]) + cuda = executor(32) + with pipeline(shape=(2, 2), executor=cuda) as p: + p.push([query]) + with p.pop() as batch: + np.testing.assert_array_equal(read_batch(batch), [[[3, 2], [1, 0]]]) + assert p.stats().gpu_bytes_committed <= cuda.limits.max_gpu_memory_bytes diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 05152a6e..d2872158 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -79,8 +79,10 @@ add_src_lib(damacy_config SOURCES damacy_config.h damacy_config.c LINKS dtype log platform ) +add_src_lib(query_selection SOURCES query/selection.h query/selection.c) + add_src_lib(lookahead SOURCES lookahead/lookahead.h lookahead/lookahead.c - LINKS log platform + LINKS log platform query_selection ) add_src_lib(prefetch_cache SOURCES prefetch/prefetch_cache.h prefetch/prefetch_cache.c @@ -147,7 +149,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 dispatch_utils assemble gpu_budget + LINKS input_transfer decoder fanout dispatch_utils execution_plan assemble gpu_budget damacy_config log CUDA::cuda_driver ) @@ -260,7 +262,9 @@ add_cuda_lib(assemble SOURCES assemble/assemble.h assemble/assemble.cu LINKS CUDA::cudart_static CUDA::cuda_driver log ) -add_src_lib(prepared_plan SOURCES planner/plan.h planner/plan.c) +add_src_lib(prepared_plan SOURCES planner/plan.h planner/plan.c + LINKS query_selection +) 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 @@ -287,6 +291,7 @@ if(NOT DAMACY_FUZZ) 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 + query_selection ) if(DAMACY_CUDA) add_cuda_lib(cuda_executor diff --git a/src/assemble/assemble.cu b/src/assemble/assemble.cu index 77ce509d..8a0fb721 100644 --- a/src/assemble/assemble.cu +++ b/src/assemble/assemble.cu @@ -199,6 +199,72 @@ src_bpe(uint8_t src_dtype) } } +__device__ __forceinline__ float +load_chunk_value(const struct sample_plan& s, + const struct assemble_chunk& c, + const uint8_t* arena_base, + uint64_t source) +{ + if (c.is_fill) + return load_src_as_float(s.fill_value, s.src_dtype); + const uint8_t* chunk_base = arena_base + c.src_base_byte_off; + switch ((enum assemble_shuffle_mode)c.shuffle_mode) { + case ASSEMBLE_SHUFFLE_BYTE: + return load_src_as_float_unshuffled(chunk_base, + source, + c.shuffle_typesize, + c.shuffle_blocksize, + s.src_dtype); + case ASSEMBLE_SHUFFLE_BIT: + return load_src_as_float_bitunshuffled(chunk_base, + source, + c.shuffle_typesize, + c.shuffle_blocksize, + s.src_dtype); + default: + return load_src_as_float(chunk_base + source * src_bpe(s.src_dtype), + s.src_dtype); + } +} + +template +__device__ __forceinline__ void +gather_body(int rank, + const struct sample_plan& s, + const struct assemble_chunk& c, + const uint8_t* arena_base, + uint8_t* output_base) +{ + const struct gather_dim* gather = s.gather_dims + c.gather_offset; + uint64_t elements = 1; + for (int d = 0; d < rank; ++d) + elements *= gather[d].count; + for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + i < elements; + i += (uint64_t)gridDim.x * blockDim.x) { + uint64_t remaining = i; + uint64_t source = 0; + uint64_t destination = s.sample_dst_off_elems; +#pragma unroll + for (int d = rank - 1; d >= 0; --d) { + const struct gather_dim dim = gather[d]; + uint32_t position = (uint32_t)(remaining % dim.count); + remaining /= dim.count; + uint32_t src = dim.begin + position; + uint64_t dst = dim.output_begin + position; + if (s.dims[d].index_count) { + const struct gather_index index = s.indices[src]; + src = index.source; + dst = index.output; + } + source += (uint64_t)src * s.dims[d].src_stride; + destination += dst * s.dims[d].dst_stride; + } + ((dst_t*)output_base)[destination] = + cast_to_dst(load_chunk_value(s, c, arena_base, source)); + } +} + // Per-thread copy. The decode of blockIdx.x → tile-in-chunk and the // per-dim accumulation of src/dst offsets and bounds fuse into one // reverse loop: the accumulators commute across d, so the row-major @@ -238,42 +304,7 @@ assemble_body(int rank, dst_t* dst = (dst_t*)(output_base + dst_off_elems * (int64_t)sizeof(dst_t)); - // Fill-mode chunks: read the broadcast value from the sample's - // fill_value (an array-level zarr property) instead of the arena. - if (c.is_fill) { - float v = load_src_as_float(s.fill_value, s.src_dtype); - *dst = cast_to_dst(v); - return; - } - - const uint32_t sbpe = src_bpe(s.src_dtype); - const uint8_t* chunk_base = arena_base + c.src_base_byte_off; - - // Each block handles one chunk, so the shuffle switch is uniform - // across the block — no warp divergence. - float v; - switch ((enum assemble_shuffle_mode)c.shuffle_mode) { - case ASSEMBLE_SHUFFLE_BYTE: - v = load_src_as_float_unshuffled(chunk_base, - src_off_elems, - c.shuffle_typesize, - c.shuffle_blocksize, - s.src_dtype); - break; - case ASSEMBLE_SHUFFLE_BIT: - v = load_src_as_float_bitunshuffled(chunk_base, - src_off_elems, - c.shuffle_typesize, - c.shuffle_blocksize, - s.src_dtype); - break; - case ASSEMBLE_SHUFFLE_NONE: - default: - v = load_src_as_float(chunk_base + src_off_elems * (int64_t)sbpe, - s.src_dtype); - break; - } - *dst = cast_to_dst(v); + *dst = cast_to_dst(load_chunk_value(s, c, arena_base, src_off_elems)); } // Single kernel template covering both the templated-rank fast path @@ -285,10 +316,13 @@ assemble_kernel(const struct sample_plan* __restrict__ d_samples, const uint8_t* __restrict__ arena_base, uint8_t* __restrict__ output_base) { - const struct assemble_chunk c = d_chunks[blockIdx.y]; + const struct assemble_chunk& c = d_chunks[blockIdx.y]; const struct sample_plan& s = d_samples[c.sample_idx_in_batch]; const int rank = (RANK_TPL == 0) ? (int)s.rank : RANK_TPL; - assemble_body(rank, s, c, arena_base, output_base); + if (s.indexed) + gather_body(rank, s, c, arena_base, output_base); + else + assemble_body(rank, s, c, arena_base, output_base); } // Dispatch kernel launch over destination dtype. RANK_TPL_VAL == 0 @@ -382,3 +416,10 @@ assemble_blocks_per_chunk(uint8_t rank, const struct sample_dim* dims) bpc *= ceil_div_u32(dims[d].chunk_shape, tile_T((int)rank, d)); return bpc; } + +extern "C" uint32_t +assemble_gather_blocks(uint64_t elements) +{ + uint64_t blocks = elements / kBlockSize + (elements % kBlockSize != 0); + return (uint32_t)(blocks > 65535 ? 65535 : blocks); +} diff --git a/src/assemble/assemble.h b/src/assemble/assemble.h index 880320c8..bfd2dc27 100644 --- a/src/assemble/assemble.h +++ b/src/assemble/assemble.h @@ -2,11 +2,13 @@ // device arena into the output batch tensor, casting source elements to // the destination dtype on the fly. // -// The kernel iterates the union of chunks for each sample (a single -// rectangle per sample, U[d] = N[d] * S[d]) and culls voxels outside -// the sample's tight AABB at write time. Per-sample constants live in -// `struct sample_plan` (planner.h); per-wave-chunk records carry the -// arena offset and the chunk's grid position. Source dtype is read off +// For a rectangular sample the kernel iterates the union of its chunks +// (U[d] = N[d] * S[d]) and culls voxels outside the sample's tight AABB at +// write time. For an indexed sample each chunk writes only the Cartesian +// product of its selected positions, described by the sample's gather_dims +// and indices. Per-sample constants live in `struct sample_plan` +// (dispatch.h); per-wave-chunk records carry the arena offset and either the +// chunk's grid position or its gather_dims offset. Source dtype is read off // the sample_plan; destination dtype is fixed for the launch. #pragma once @@ -46,6 +48,7 @@ extern "C" uint8_t shuffle_mode; // enum assemble_shuffle_mode uint8_t shuffle_typesize; // 1, 2, 4, 8; 0 if NONE uint8_t is_fill; // 1 = broadcast sample.fill_value + uint32_t gather_offset; // index into sample.gather_dims[] uint32_t chunk_d[DAMACY_MAX_RANK]; // chunk grid position within sample }; @@ -84,6 +87,8 @@ extern "C" // Compute blocks-per-chunk for a sample given its dims and the kernel's // block-tile shape. Host-side helper used by wave orchestration when // packing the wave's metadata. Returns 0 for unsupported ranks. + uint32_t assemble_gather_blocks(uint64_t elements); + uint32_t assemble_blocks_per_chunk(uint8_t rank, const struct sample_dim* dims); diff --git a/src/damacy.h b/src/damacy.h index 5b677c21..60596987 100644 --- a/src/damacy.h +++ b/src/damacy.h @@ -74,11 +74,35 @@ extern "C" uint8_t rank; }; - // One sample request. + struct damacy_index_array + { + const int64_t* values; + uint32_t count; + }; + + enum damacy_axis_kind + { + DAMACY_AXIS_INTERVAL = 1, + DAMACY_AXIS_INDICES = 2, + }; + + struct damacy_axis_selection + { + enum damacy_axis_kind kind; + union + { + struct damacy_interval interval; + struct damacy_index_array indices; + }; + }; + + // Axes combine as a Cartesian product, preserving index order and duplicates. + // Accepted samples own a copy of the URI and index values. struct damacy_sample { - const char* uri; // null-terminated; copied internally - struct damacy_aabb aabb; + const char* uri; + uint8_t rank; + struct damacy_axis_selection axes[DAMACY_MAX_RANK]; }; // Caller-owned slice of samples to push. @@ -96,6 +120,11 @@ extern "C" // Required — no default; a value too small for the requested // geometry returns DAMACY_BUDGET from damacy_create. uint64_t max_gpu_memory_bytes; + // Index storage per batch, 8 bytes per index, counted against + // max_gpu_memory_bytes. Must be 0, which rejects indexed samples at + // push, or in [8 * samples_per_batch * sum(sample_shape), 8 * UINT32_MAX] + // so a batch can never run out. + uint64_t max_index_bytes; // Required: must be in [1, DAMACY_MAX_CHUNK_BYTES]. uint32_t max_chunk_uncompressed_bytes; // Required: must be in [1, UINT32_MAX] (read_op.nbytes is uint32_t). @@ -169,7 +198,7 @@ extern "C" // Destination dtype of assembled batches. enum damacy_dtype dtype; // Per-sample output extents (in voxels) along the zarr's axis order — same - // layout damacy_sample.aabb uses. sample_rank must be in + // axis order damacy_sample.axes uses. sample_rank must be in // [1, DAMACY_MAX_RANK]; every sample_shape[d] must be > 0. int64_t sample_shape[DAMACY_MAX_RANK]; uint8_t sample_rank; @@ -243,6 +272,8 @@ extern "C" // a batch (or wait) and retry with the returned suffix // INVAL bad arguments (samples.beg > samples.end, null d, etc.) // RANK sample rank incompatible with cfg.sample_rank + // BUDGET the sample's index arrays exceed the plan storage limit, or + // it has index arrays and tuning.max_index_bytes is 0 // SHUTDOWN instance is in a failed state or being destroyed // Store-derived errors (missing uri, unsupported source dtype, per-array // rank mismatch, decode failures) surface asynchronously from damacy_pop. diff --git a/src/damacy_config.c b/src/damacy_config.c index 293cc5ed..6408e8c4 100644 --- a/src/damacy_config.c +++ b/src/damacy_config.c @@ -48,6 +48,8 @@ validate_config(const struct damacy_config* cfg) CHECK_SILENT(Invalid, cfg); CHECK_SILENT(Invalid, cfg->samples_per_batch > 0); CHECK_SILENT(Invalid, cfg->tuning.max_gpu_memory_bytes > 0); + CHECK_SILENT(Invalid, + cfg->tuning.max_index_bytes <= UINT64_C(8) * UINT32_MAX); CHECK_SILENT(Invalid, cfg->lookahead_samples >= cfg->samples_per_batch); CHECK_SILENT(Invalid, cfg->tuning.n_io_threads > 0); CHECK_SILENT(Invalid, cfg->tuning.n_io_threads <= DAMACY_MAX_IO_THREADS); @@ -107,8 +109,8 @@ validate_config(const struct damacy_config* cfg) // name the knob, observed vs required value, and the fix. See // dev/metadata_prefetch.md. { - uint64_t meta_floor = - (uint64_t)cfg->lookahead_samples + 2ull * (uint64_t)cfg->samples_per_batch; + uint64_t meta_floor = (uint64_t)cfg->lookahead_samples + + 2ull * (uint64_t)cfg->samples_per_batch; if ((uint64_t)cfg->tuning.n_array_meta_cache < meta_floor) { log_error("n_array_meta_cache=%u is too small: requires >= " "lookahead_samples(%u) + 2*samples_per_batch(%u) = %llu. Raise " @@ -183,6 +185,7 @@ damacy_tuning_defaults(void) return (struct damacy_tuning){ .max_chunk_uncompressed_bytes = 2 * 1024 * 1024, .max_read_op_bytes = 4 * 1024 * 1024, + .max_index_bytes = 64ull << 20, .host_buffer_waves = DAMACY_DEFAULT_HOST_BUFFER_WAVES, .max_chunks_per_wave = DAMACY_DEFAULT_MAX_CHUNKS_PER_WAVE, .max_substreams_per_chunk = DAMACY_DEFAULT_MAX_SUBSTREAMS_PER_CHUNK, diff --git a/src/damacy_lifecycle.c b/src/damacy_lifecycle.c index 2fd8e5a4..4ee7151f 100644 --- a/src/damacy_lifecycle.c +++ b/src/damacy_lifecycle.c @@ -209,6 +209,7 @@ damacy_create(const struct damacy_config* config, struct damacy** out) &(struct damacy_cuda_config){ .device = config->device, .max_gpu_memory_bytes = tuning->max_gpu_memory_bytes, + .max_index_bytes = tuning->max_index_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, diff --git a/src/damacy_pipeline.h b/src/damacy_pipeline.h index 7f09c6cf..c6765736 100644 --- a/src/damacy_pipeline.h +++ b/src/damacy_pipeline.h @@ -54,6 +54,7 @@ extern "C" { int device; uint64_t max_gpu_memory_bytes; + uint64_t max_index_bytes; uint32_t max_chunk_bytes; uint64_t max_read_bytes; uint32_t max_chunks_per_wave; diff --git a/src/damacy_push.c b/src/damacy_push.c index 9f79e8f8..cf89d0e2 100644 --- a/src/damacy_push.c +++ b/src/damacy_push.c @@ -1,5 +1,24 @@ #include "damacy_internal.h" +#include "query/selection.h" + +static struct damacy_push_result +push_to_planner(struct damacy* self, struct damacy_sample_slice samples) +{ + struct damacy_sample_slice accepted = samples; + if (!self->executor->accepts_indexed_samples) { + accepted.end = samples.beg; + while (accepted.end != samples.end && !query_has_indices(accepted.end)) + ++accepted.end; + } + struct damacy_push_result result = + self->planner->ops->push(self->planner, accepted); + result.unconsumed.end = samples.end; + if (result.status == DAMACY_OK && accepted.end != samples.end) + result.status = DAMACY_BUDGET; + return result; +} + struct damacy_push_result damacy_push(struct damacy* self, struct damacy_sample_slice samples) { @@ -16,7 +35,7 @@ damacy_push(struct damacy* self, struct damacy_sample_slice samples) if (self->stopping || self->failed_status != DAMACY_OK) result.status = DAMACY_SHUTDOWN; else - result = self->planner->ops->push(self->planner, samples); + result = push_to_planner(self, samples); scheduler_broadcast(self->sched); scheduler_unlock(self->sched); return result; diff --git a/src/executor/cpu_executor.c b/src/executor/cpu_executor.c index 897b5fbf..820d8606 100644 --- a/src/executor/cpu_executor.c +++ b/src/executor/cpu_executor.c @@ -242,6 +242,73 @@ source_value(const void* source, enum dtype dtype, uint64_t index) } } +static uint64_t +gather_chunk(struct cpu_executor* self, + struct cpu_slot* slot, + const struct plan_chunk* chunk, + const struct plan_region* region, + const struct zarr_metadata* meta, + const void* decoded, + float fill) +{ + struct selection_span spans[DAMACY_MAX_RANK]; + uint64_t origin[DAMACY_MAX_RANK]; + uint64_t position[DAMACY_MAX_RANK] = { 0 }; + uint64_t elements = 1; + for (uint8_t d = 0; d < meta->rank; ++d) { + origin[d] = chunk->coordinate[d] * meta->inner_chunk_shape[d]; + spans[d] = query_chunk_span(®ion->axes[d], + region->source.dims[d], + origin[d], + meta->inner_chunk_shape[d]); + elements *= spans[d].count; + } + if (!elements) + return 0; + uint8_t last = meta->rank - 1; + uint64_t width = region->axes[last].count ? 1 : spans[last].count; + for (;;) { + uint64_t source = 0; + uint64_t destination = (uint64_t)region->sample * self->strides[0]; + for (uint8_t d = 0; d < meta->rank; ++d) { + const struct query_axis* axis = ®ion->axes[d]; + uint64_t p = spans[d].begin + position[d]; + uint64_t coordinate = axis->count ? (uint64_t)axis->indices[p].source : p; + uint64_t output = axis->count + ? axis->indices[p].output + : coordinate - (uint64_t)region->source.dims[d].beg; + source = source * meta->inner_chunk_shape[d] + coordinate - origin[d]; + destination += output * (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 { + for (uint64_t j = 0; j < width; ++j) { + float value = chunk->missing + ? fill + : source_value(decoded, meta->dtype, source + j); + if (self->output.dtype == DAMACY_F32) + ((float*)slot->buffer->data)[destination + j] = value; + else + ((uint16_t*)slot->buffer->data)[destination + j] = + float_to_bfloat(value); + } + } + int d = region->axes[last].count ? last : (int)last - 1; + for (; d >= 0; --d) { + if (++position[d] < spans[d].count) + break; + position[d] = 0; + } + if (d < 0) + break; + } + return elements * damacy_dtype_bpe(self->output.dtype); +} + static uint64_t assemble_chunk(struct cpu_executor* self, struct cpu_slot* slot, @@ -254,6 +321,11 @@ assemble_chunk(struct cpu_executor* self, 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]; + if (region->operation == PLAN_GATHER) { + output_bytes += + gather_chunk(self, slot, chunk, region, meta, decoded, fill); + continue; + } 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); @@ -778,6 +850,7 @@ damacy_cpu_executor_create(struct damacy_reader* reader, return DAMACY_OOM; self->base.ops = &cpu_ops; self->base.device_type = DAMACY_DEVICE_CPU; + self->base.accepts_indexed_samples = 1; self->reader = reader; self->config = *config; *out = &self->base; diff --git a/src/executor/cuda_executor.c b/src/executor/cuda_executor.c index ac20f70b..e54f8542 100644 --- a/src/executor/cuda_executor.c +++ b/src/executor/cuda_executor.c @@ -151,6 +151,7 @@ cuda_start(struct damacy_executor* base, .samples_per_batch = output->samples_per_batch, .device = self->config.device, .tuning = { .max_gpu_memory_bytes = self->config.max_gpu_memory_bytes, + .max_index_bytes = self->config.max_index_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, @@ -167,6 +168,18 @@ cuda_start(struct damacy_executor* base, memcpy(self->cfg.sample_shape, output->sample_shape, sizeof(self->cfg.sample_shape)); + uint64_t max_indices = dispatch_max_indices(&self->cfg); + if (self->config.max_index_bytes && + self->config.max_index_bytes / sizeof(struct gather_index) < + max_indices) { + log_error("max_index_bytes=%llu is too small: a batch can hold %llu " + "indices (samples_per_batch * sum(sample_shape)), 8 bytes each. " + "Raise max_index_bytes, or set it to 0 to disable indexed " + "queries", + (unsigned long long)self->config.max_index_bytes, + (unsigned long long)max_indices); + return DAMACY_BUDGET; + } enum damacy_status status = DAMACY_CUDA; int pushed = 0; if (cuInit(0) != CUDA_SUCCESS) @@ -247,7 +260,10 @@ cuda_start(struct damacy_executor* base, 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); + failed = render_job_init(&self->jobs.jobs[i], + output->samples_per_batch, + dispatch_index_capacity(&self->cfg), + dispatch_gather_dim_capacity(&self->cfg)); if (!failed) failed = wave_pool_init(&self->waves, &self->batches, @@ -580,6 +596,8 @@ damacy_cuda_executor_create(struct damacy_reader* reader, *out = NULL; if (!reader || !config || config->device < -1 || !config->max_gpu_memory_bytes || !config->chunk_layout_entries || + config->max_index_bytes > + (uint64_t)UINT32_MAX * sizeof(struct gather_index) || !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 || @@ -598,6 +616,7 @@ damacy_cuda_executor_create(struct damacy_reader* reader, return DAMACY_OOM; self->base.ops = &cuda_ops; self->base.device_type = DAMACY_DEVICE_CUDA; + self->base.accepts_indexed_samples = config->max_index_bytes != 0; self->reader = reader; self->config = *config; *out = &self->base; diff --git a/src/executor/cuda_geometry.inc b/src/executor/cuda_geometry.inc index 4ef3b40f..e0057e65 100644 --- a/src/executor/cuda_geometry.inc +++ b/src/executor/cuda_geometry.inc @@ -30,6 +30,11 @@ resolve_wave_geometry(const struct damacy_config* cfg, out->input = out->want_gds ? input_transfer_gds() : input_transfer_host_staging(); + uint64_t index_bytes = 2ull * dispatch_index_storage_bytes(cfg); + if (index_bytes > resolver_budget) + return DAMACY_BUDGET; + resolver_budget -= index_bytes; + const struct input_transfer_resources min_input = input_transfer_resources(out->input, out->host_buffer_waves, 0); enum damacy_status s = @@ -47,6 +52,7 @@ resolve_wave_geometry(const struct damacy_config* cfg, return s; } + out->sizing.worst_case_total_bytes += index_bytes; const struct input_transfer_resources input_resources = input_transfer_resources( out->input, out->host_buffer_waves, out->sizing.input_staging_per_wave); diff --git a/src/executor/dispatch.c b/src/executor/dispatch.c index 52684c0e..e2febf6c 100644 --- a/src/executor/dispatch.c +++ b/src/executor/dispatch.c @@ -33,6 +33,55 @@ scratch_reserve(struct dispatch_scratch* scratch, uint32_t count) return DAMACY_OK; } +uint64_t +dispatch_max_indices(const struct damacy_config* config) +{ + if (!config || !config->samples_per_batch || !config->sample_rank || + config->sample_rank > DAMACY_MAX_RANK) + return 0; + uint64_t count = 0; + for (uint8_t d = 0; d < config->sample_rank; ++d) { + if (config->sample_shape[d] <= 0) + return 0; + uint64_t extent = (uint64_t)config->sample_shape[d]; + if (extent > UINT64_MAX - count) + return UINT64_MAX; + count += extent; + } + if (count > UINT64_MAX / config->samples_per_batch) + return UINT64_MAX; + return count * config->samples_per_batch; +} + +uint32_t +dispatch_index_capacity(const struct damacy_config* config) +{ + uint64_t capacity = dispatch_max_indices(config); + if (!capacity) + return 0; + uint64_t limit = config->tuning.max_index_bytes / sizeof(struct gather_index); + if (capacity > limit) + capacity = limit; + return capacity > UINT32_MAX ? UINT32_MAX : (uint32_t)capacity; +} + +uint32_t +dispatch_gather_dim_capacity(const struct damacy_config* config) +{ + return dispatch_index_capacity(config) + ? DAMACY_MAX_CHUNKS_PER_BATCH * config->sample_rank + : 0; +} + +uint64_t +dispatch_index_storage_bytes(const struct damacy_config* config) +{ + return (uint64_t)dispatch_index_capacity(config) * + sizeof(struct gather_index) + + (uint64_t)dispatch_gather_dim_capacity(config) * + sizeof(struct gather_dim); +} + enum damacy_status dispatch_plan_build(const struct prepared_plan* plan, uint16_t slot, @@ -50,6 +99,7 @@ dispatch_plan_build(const struct prepared_plan* plan, 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; + out->n_indices = out->n_gather_dims = 0; int64_t strides[DAMACY_MAX_RANK + 1]; strides[plan->output.sample_rank] = 1; for (int d = plan->output.sample_rank - 1; d >= 0; --d) @@ -63,25 +113,37 @@ dispatch_plan_build(const struct prepared_plan* plan, .sample_idx_in_batch = (uint16_t)region->sample, .rank = meta->rank, .src_dtype = (uint8_t)meta->dtype, + .indexed = region->operation == PLAN_GATHER, .sample_dst_off_elems = - (int64_t)region->sample * strides[0], - .chunk_count = 1 }; + (int64_t)region->sample * strides[0] }; 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); + const struct query_axis* axis = ®ion->axes[d]; + if (axis->count > out->indices_cap - out->n_indices) + return DAMACY_BUDGET; + if (axis->count && !out->indices) + return DAMACY_INVAL; + uint32_t index_offset = out->n_indices; + for (uint32_t j = 0; j < axis->count; ++j) + out->indices[out->n_indices++] = (struct gather_index){ + .source = (uint32_t)((uint64_t)axis->indices[j].source % chunk), + .output = axis->indices[j].output + }; + sample->dims[d] = + (struct sample_dim){ .chunk_shape = (uint32_t)chunk, + .chunk_grid_extent = + axis->count ? 0 : (uint32_t)(end - begin), + .index_offset = index_offset, + .index_count = axis->count, + .aabb_lo_relative = region->source.dims[d].beg - + (int64_t)(begin * chunk), + .aabb_extent = plan->output.sample_shape[d], + .dst_stride = strides[d + 1], + .src_stride = source_stride }; source_stride *= (int64_t)chunk; } } @@ -104,10 +166,35 @@ dispatch_plan_build(const struct prepared_plan* plan, .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 / + struct sample_plan* sample = &out->sample_plans[region->sample]; + if (sample->indexed) { + if (meta->rank > out->gather_dims_cap - out->n_gather_dims) + return DAMACY_BUDGET; + dispatch->gather_offset = out->n_gather_dims; + out->n_gather_dims += meta->rank; + } + for (uint8_t d = 0; d < meta->rank; ++d) { + if (!sample->indexed) { + dispatch->chunk_d[d] = (uint32_t)(chunk->coordinate[d] - + (uint64_t)region->source.dims[d].beg / meta->inner_chunk_shape[d]); + continue; + } + uint64_t origin = chunk->coordinate[d] * meta->inner_chunk_shape[d]; + struct selection_span span = query_chunk_span(®ion->axes[d], + region->source.dims[d], + origin, + meta->inner_chunk_shape[d]); + out->gather_dims[dispatch->gather_offset + d] = (struct gather_dim){ + .begin = region->axes[d].count + ? sample->dims[d].index_offset + (uint32_t)span.begin + : (uint32_t)(span.begin - origin), + .count = (uint32_t)span.count, + .output_begin = region->axes[d].count + ? 0 + : span.begin - (uint64_t)region->source.dims[d].beg + }; + } if (!chunk->missing) { uint64_t start = chunk->offset / alignment * alignment; uint64_t end = chunk->offset + chunk->encoded_bytes; diff --git a/src/executor/dispatch.h b/src/executor/dispatch.h index 976cd673..242f33dd 100644 --- a/src/executor/dispatch.h +++ b/src/executor/dispatch.h @@ -18,10 +18,25 @@ extern "C" uint32_t nbytes; }; + struct gather_index + { + uint32_t source; + uint32_t output; + }; + + struct gather_dim + { + uint32_t begin; + uint32_t count; + uint64_t output_begin; + }; + struct sample_dim { uint32_t chunk_shape; uint32_t chunk_grid_extent; + uint32_t index_offset; + uint32_t index_count; int64_t aabb_lo_relative; int64_t aabb_extent; int64_t dst_stride; @@ -34,11 +49,12 @@ extern "C" uint16_t sample_idx_in_batch; uint8_t rank; uint8_t src_dtype; + uint8_t indexed; + const struct gather_index* indices; + const struct gather_dim* gather_dims; 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; @@ -53,6 +69,7 @@ extern "C" uint32_t decompressed_nbytes; uint64_t host_buf_offset; uint32_t dev_decompressed_offset; + uint32_t gather_offset; uint16_t batch_pool_slot; uint16_t sample_idx_in_batch; uint8_t codec_id; @@ -96,6 +113,12 @@ extern "C" struct read_op_group* read_op_groups; uint32_t read_op_groups_cap; uint32_t n_read_op_groups; + struct gather_index* indices; + uint32_t indices_cap; + uint32_t n_indices; + struct gather_dim* gather_dims; + uint32_t gather_dims_cap; + uint32_t n_gather_dims; struct path_intern* paths; uint32_t n_chunks_to_load; uint32_t n_loads_issued; @@ -109,6 +132,11 @@ extern "C" uint32_t capacity; }; + uint64_t dispatch_max_indices(const struct damacy_config* config); + uint32_t dispatch_index_capacity(const struct damacy_config* config); + uint32_t dispatch_gather_dim_capacity(const struct damacy_config* config); + uint64_t dispatch_index_storage_bytes(const struct damacy_config* config); + enum damacy_status dispatch_plan_build(const struct prepared_plan* plan, uint16_t slot, uint64_t alignment, diff --git a/src/lookahead/lookahead.c b/src/lookahead/lookahead.c index 0623f614..b76afb53 100644 --- a/src/lookahead/lookahead.c +++ b/src/lookahead/lookahead.c @@ -13,9 +13,7 @@ sample_slot_clear(struct damacy_sample_slot* slot) if (!slot) return; free(slot->uri); - slot->uri = NULL; - memset(&slot->aabb, 0, sizeof(slot->aabb)); - slot->sample_seq = 0; + *slot = (struct damacy_sample_slot){ 0 }; } int @@ -73,10 +71,8 @@ lookahead_push_with_sample_seq(struct damacy_lookahead* la, if (la->size == la->cap) goto Unlock; struct damacy_sample_slot* slot = &la->slots[la->tail]; - slot->uri = strdup(sample->uri); - if (!slot->uri) + if (query_copy(sample, &slot->uri, &slot->aabb, slot->axes) != DAMACY_OK) goto Unlock; - slot->aabb = sample->aabb; slot->sample_seq = sample_seq; la->tail = (la->tail + 1) % la->cap; la->size++; diff --git a/src/lookahead/lookahead.h b/src/lookahead/lookahead.h index 28ae1da6..6911a2ea 100644 --- a/src/lookahead/lookahead.h +++ b/src/lookahead/lookahead.h @@ -1,5 +1,7 @@ #pragma once +#include "query/selection.h" + #include "damacy.h" #include @@ -11,6 +13,7 @@ struct damacy_sample_slot { char* uri; struct damacy_aabb aabb; + struct query_axis axes[DAMACY_MAX_RANK]; uint64_t sample_seq; }; diff --git a/src/pipeline/components.h b/src/pipeline/components.h index 6838f3aa..890550e3 100644 --- a/src/pipeline/components.h +++ b/src/pipeline/components.h @@ -71,6 +71,7 @@ struct damacy_executor _Atomic int active; int device_type; int device_id; + int accepts_indexed_samples; }; struct damacy_buffer diff --git a/src/pipeline/zarr_planner.c b/src/pipeline/zarr_planner.c index 840c5bfb..eb210945 100644 --- a/src/pipeline/zarr_planner.c +++ b/src/pipeline/zarr_planner.c @@ -9,6 +9,7 @@ #include "store/metadata_store_async.h" #include +#include struct zarr_planner { @@ -138,22 +139,10 @@ zarr_push(struct damacy_planner* base, struct damacy_sample_slice samples) result.status = DAMACY_AGAIN; return result; } - if (!sample->uri) { - result.status = DAMACY_INVAL; + result.status = + query_validate(sample, &self->output, self->limits.max_plan_bytes); + if (result.status != DAMACY_OK) 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; @@ -186,6 +175,8 @@ zarr_next(struct damacy_planner* base, struct prepared_plan** out) .h_meta = ready.h_meta, .h_shards = ready.h_shards, .n_shards = ready.n_shards }; + memcpy( + self->samples[self->staged - 1].axes, ready.axes, sizeof(ready.axes)); self->watermark = ready.sample_seq + 1; } enum damacy_status status = prepared_plan_build(self->arrays, diff --git a/src/planner/plan.h b/src/planner/plan.h index 2b41b2b0..2fa8d43c 100644 --- a/src/planner/plan.h +++ b/src/planner/plan.h @@ -1,6 +1,7 @@ #pragma once #include "damacy_pipeline.h" +#include "query/selection.h" #include "zarr/zarr_metadata.h" struct plan_array @@ -24,6 +25,7 @@ struct plan_chunk enum plan_operation { PLAN_COPY, + PLAN_GATHER, }; struct plan_region @@ -32,6 +34,7 @@ struct plan_region uint32_t array; uint32_t sample; struct damacy_aabb source; + struct query_axis axes[DAMACY_MAX_RANK]; }; struct plan_use diff --git a/src/planner/plan_builder.c b/src/planner/plan_builder.c index 7d7d8f35..15124cdc 100644 --- a/src/planner/plan_builder.c +++ b/src/planner/plan_builder.c @@ -5,7 +5,6 @@ #include "prefetch/shard_index.h" #include "util/hash.h" #include "util/strbuf.h" -#include "zarr/sample_shard_iterator.h" #include #include @@ -32,8 +31,6 @@ 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) { @@ -41,7 +38,6 @@ sample_geometry(const struct planner_sample* sample, 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; @@ -49,18 +45,24 @@ sample_geometry(const struct planner_sample* sample, 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]) + int64_t extent = sample->axes[d].count ? sample->axes[d].count : hi - lo; + if (extent != 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; + struct selection_grid grid; + enum damacy_status status = selection_grid_init(&grid, + &sample->aabb, + sample->axes, + meta->inner_chunk_shape, + NULL, + NULL, + limits->max_chunks); + if (status != DAMACY_OK) + return status; + *count = (uint32_t)grid.count; *bytes = (uint32_t)size; return DAMACY_OK; } @@ -105,6 +107,7 @@ prepared_plan_build(struct prefetch_cache* arrays, struct strbuf path = { 0 }; uint64_t capacity = 0; uint64_t path_bytes = 0; + uint64_t index_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) @@ -112,12 +115,18 @@ prepared_plan_build(struct prefetch_cache* arrays, 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); + status = sample_geometry(&samples[i], meta, output, limits, &count, &bytes); if (status != DAMACY_OK) return status; + for (uint8_t d = 0; d < meta->rank; ++d) { + uint64_t size = + (uint64_t)samples[i].axes[d].count * sizeof(struct query_index); + if (size > limits->max_plan_bytes || + index_bytes > limits->max_plan_bytes - size) + return DAMACY_BUDGET; + index_bytes += size; + } capacity += count; if (capacity > limits->max_chunks) return DAMACY_BUDGET; @@ -140,6 +149,9 @@ prepared_plan_build(struct prefetch_cache* arrays, (sizeof(struct plan_array) + sizeof(struct plan_region)) + capacity * (sizeof(struct plan_chunk) + sizeof(struct plan_use)) + buckets * sizeof(uint32_t) + path_bytes; + if (index_bytes > UINT64_MAX - storage_bytes) + return DAMACY_BUDGET; + storage_bytes += index_bytes; if (storage_bytes > SIZE_MAX || storage_bytes > limits->max_plan_bytes || sizeof(*plan) > limits->max_plan_bytes - storage_bytes) return DAMACY_BUDGET; @@ -156,7 +168,8 @@ prepared_plan_build(struct prefetch_cache* arrays, plan->arrays = plan->storage; plan->regions = (void*)(plan->arrays + n_samples); plan->chunks = (void*)(plan->regions + n_samples); - plan->uses = (void*)(plan->chunks + capacity); + struct query_index* indices = (void*)(plan->chunks + capacity); + plan->uses = (void*)((char*)indices + index_bytes); uint32_t* table = (void*)(plan->uses + capacity); char* paths = (void*)(table + buckets); for (uint32_t i = 0; i < n_samples; ++i) { @@ -178,26 +191,39 @@ prepared_plan_build(struct prefetch_cache* arrays, .array = array, .sample = i, .source = sample->aabb }; + for (uint8_t d = 0; d < meta->rank; ++d) { + uint32_t n = sample->axes[d].count; + if (!n) + continue; + memcpy(indices, sample->axes[d].indices, (size_t)n * sizeof(*indices)); + plan->regions[i].axes[d] = (struct query_axis){ indices, n }; + plan->regions[i].operation = PLAN_GATHER; + indices += n; + } ++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); + status = + sample_geometry(sample, meta, output, limits, &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; + struct selection_grid iterator; + status = selection_grid_init(&iterator, + &sample->aabb, + sample->axes, + meta->shard_shape, + NULL, + NULL, + limits->max_shards_per_sample); + if (status != DAMACY_OK) goto Done; - } uint64_t shard_coord[DAMACY_MAX_RANK]; uint32_t shard_index = 0; - while (sample_shard_iterator_next(&iterator, shard_coord)) { + while (selection_grid_next(&iterator, shard_coord)) { if (shard_index >= sample->n_shards) { status = DAMACY_INVAL; goto Done; @@ -226,13 +252,20 @@ prepared_plan_build(struct prefetch_cache* arrays, 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]; + lo[d] = shard_coord[d] * per_shard[d]; + hi[d] = lo[d] + per_shard[d]; } - for (;;) { + struct selection_grid chunks; + status = selection_grid_init(&chunks, + &sample->aabb, + sample->axes, + meta->inner_chunk_shape, + lo, + hi, + limits->max_chunks); + if (status != DAMACY_OK) + goto Done; + while (selection_grid_next(&chunks, coordinate)) { uint64_t entry_index = 0; for (uint8_t d = 0; d < meta->rank; ++d) entry_index = @@ -289,16 +322,6 @@ prepared_plan_build(struct prefetch_cache* arrays, .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) { diff --git a/src/planner/plan_builder.h b/src/planner/plan_builder.h index f7da85d0..c36a3de7 100644 --- a/src/planner/plan_builder.h +++ b/src/planner/plan_builder.h @@ -1,5 +1,7 @@ #pragma once +#include "query/selection.h" + #include "planner/plan.h" #include "prefetch/prefetch_handle.h" @@ -7,6 +9,7 @@ struct planner_sample { const char* uri; struct damacy_aabb aabb; + struct query_axis axes[DAMACY_MAX_RANK]; struct prefetch_handle h_meta; struct prefetch_handle* h_shards; uint32_t n_shards; diff --git a/src/planner/planner.c b/src/planner/planner.c index aee32e36..2efb5f2c 100644 --- a/src/planner/planner.c +++ b/src/planner/planner.c @@ -64,7 +64,8 @@ planner_plan_segment(struct planner* self, int64_t hi = samples[0].aabb.dims[d].end; if (lo < 0 || hi <= lo) return DAMACY_INVAL; - output.sample_shape[d] = hi - lo; + output.sample_shape[d] = + samples[0].axes[d].count ? samples[0].axes[d].count : hi - lo; } struct damacy_plan_limits limits = { .max_chunks = out->chunk_plans_cap < out->read_ops_cap diff --git a/src/prefetch/prefetcher.c b/src/prefetch/prefetcher.c index b2a25448..89486f4a 100644 --- a/src/prefetch/prefetcher.c +++ b/src/prefetch/prefetcher.c @@ -50,6 +50,7 @@ struct prefetcher_slot int err_code; char* uri; struct damacy_aabb aabb; + struct query_axis axes[DAMACY_MAX_RANK]; uint64_t sample_seq; struct prefetch_gate* gate; struct prefetch_handle h_meta; @@ -331,17 +332,15 @@ advance_from_meta(struct prefetcher* p, struct prefetcher_slot* s) const struct zarr_metadata* meta = (const struct zarr_metadata*)value; CHECK(Bad, meta); - struct sample_shard_iterator it; - CHECK(Bad, sample_shard_iterator_init(&it, meta, &s->aabb) == 0); - - uint64_t n = 1; - err = DAMACY_INVAL; // a shard-count overflow below fails the sample as - // an oversized (invalid) sample - for (uint8_t d = 0; d < it.rank; ++d) { - uint64_t span = it.shard_end[d] - it.shard_beg[d]; - CHECK_MUL_OVERFLOW(Bad, n, span, UINT64_MAX); - n *= span; - } + struct sample_shard_iterator bounds; + err = DAMACY_INVAL; + CHECK(Bad, sample_shard_iterator_init(&bounds, meta, &s->aabb) == 0); + struct selection_grid it; + CHECK(Bad, + selection_grid_init( + &it, &s->aabb, s->axes, meta->shard_shape, NULL, NULL, UINT64_MAX) == + DAMACY_OK); + uint64_t n = it.count; err = 0; // Per-sample shard cap. The shard_index cache is sized at config time @@ -398,8 +397,8 @@ advance_from_meta(struct prefetcher* p, struct prefetcher_slot* s) // Enumerate every coord up front so the request loop is resumable on AGAIN // without re-running the iterator. for (uint32_t i = 0; - i < n && sample_shard_iterator_next( - &it, &s->shard_coords[(size_t)i * meta->rank]); + i < n && + selection_grid_next(&it, &s->shard_coords[(size_t)i * meta->rank]); ++i) ; s->n_shards = (uint32_t)n; @@ -513,6 +512,7 @@ emit_error_slot_locked(struct prefetcher* p, .sample_seq = popped->sample_seq, .gate = gate, }; + memcpy(slot->axes, popped->axes, sizeof(slot->axes)); slot_fail(p, slot, err_code); } @@ -558,6 +558,7 @@ admit_locked(struct prefetcher* p, .gate = gate, .h_meta = req.handle, }; + memcpy(slot->axes, popped->axes, sizeof(slot->axes)); slot_set_state(slot, &state_pending_meta); p->submitted++; return DAMACY_OK; @@ -772,6 +773,7 @@ pop_terminal_slot_idx_locked(struct prefetcher* self, .n_shards = s->n_shards, .h_layout = s->h_layout, }; + memcpy(out->axes, s->axes, sizeof(out->axes)); uint64_t owner_id = s->sample_seq; free(s->shard_coords); slot_reset_free(s); diff --git a/src/prefetch/prefetcher.h b/src/prefetch/prefetcher.h index 8623a90f..bfbda103 100644 --- a/src/prefetch/prefetcher.h +++ b/src/prefetch/prefetcher.h @@ -1,5 +1,7 @@ #pragma once +#include "query/selection.h" + #include "damacy.h" #include "prefetch/prefetch_cache.h" @@ -61,6 +63,7 @@ extern "C" int err_code; // damacy_status when result == PREFETCHER_RESULT_ERROR char* uri; struct damacy_aabb aabb; + struct query_axis axes[DAMACY_MAX_RANK]; uint64_t sample_seq; struct prefetch_handle h_meta; // On ERROR mid-shard-allocation, h_shards may be non-NULL with n_shards diff --git a/src/query/selection.c b/src/query/selection.c new file mode 100644 index 00000000..f0d4d8f8 --- /dev/null +++ b/src/query/selection.c @@ -0,0 +1,269 @@ +#include "query/selection.h" + +#include +#include + +static enum damacy_status +axis_length(const struct damacy_axis_selection* axis, uint64_t* length) +{ + switch (axis->kind) { + case DAMACY_AXIS_INTERVAL: + if (axis->interval.beg < 0 || axis->interval.end <= axis->interval.beg) + return DAMACY_INVAL; + *length = (uint64_t)(axis->interval.end - axis->interval.beg); + return DAMACY_OK; + case DAMACY_AXIS_INDICES: + if (!axis->indices.values || !axis->indices.count) + return DAMACY_INVAL; + *length = axis->indices.count; + return DAMACY_OK; + default: + return DAMACY_INVAL; + } +} + +enum damacy_status +query_validate(const struct damacy_sample* sample, + const struct damacy_batch_spec* output, + uint64_t max_plan_bytes) +{ + if (!sample || !sample->uri || !output) + return DAMACY_INVAL; + if (!sample->rank || sample->rank > DAMACY_MAX_RANK || + sample->rank != output->sample_rank) + return DAMACY_RANK; + uint64_t bytes = 0; + for (uint8_t d = 0; d < sample->rank; ++d) { + const struct damacy_axis_selection* axis = &sample->axes[d]; + uint64_t length; + enum damacy_status status = axis_length(axis, &length); + if (status != DAMACY_OK) + return status; + if (length != (uint64_t)output->sample_shape[d]) + return DAMACY_INVAL; + if (axis->kind == DAMACY_AXIS_INDICES) { + uint64_t size = + (uint64_t)axis->indices.count * sizeof(struct query_index); + if (size > max_plan_bytes || bytes > max_plan_bytes - size) + return DAMACY_BUDGET; + bytes += size; + for (uint32_t i = 0; i < axis->indices.count; ++i) + if (axis->indices.values[i] < 0 || axis->indices.values[i] == INT64_MAX) + return DAMACY_INVAL; + } + } + return DAMACY_OK; +} + +int +query_has_indices(const struct damacy_sample* sample) +{ + for (uint8_t d = 0; d < sample->rank && d < DAMACY_MAX_RANK; ++d) + if (sample->axes[d].kind == DAMACY_AXIS_INDICES) + return 1; + return 0; +} + +static int +compare_indices(const void* left, const void* right) +{ + const struct query_index* a = left; + const struct query_index* b = right; + if (a->source != b->source) + return (a->source > b->source) - (a->source < b->source); + return (a->output > b->output) - (a->output < b->output); +} + +enum damacy_status +query_copy(const struct damacy_sample* sample, + char** uri, + struct damacy_aabb* bounds, + struct query_axis* axes) +{ + if (!sample || !sample->uri || !sample->rank || + sample->rank > DAMACY_MAX_RANK) + return DAMACY_INVAL; + size_t uri_bytes = strlen(sample->uri) + 1; + size_t alignment = _Alignof(struct query_index); + if (uri_bytes > SIZE_MAX - alignment + 1) + return DAMACY_BUDGET; + size_t offset = (uri_bytes + alignment - 1) / alignment * alignment; + uint64_t count = 0; + for (uint8_t d = 0; d < sample->rank; ++d) { + uint64_t length; + enum damacy_status status = axis_length(&sample->axes[d], &length); + if (status != DAMACY_OK) + return status; + if (sample->axes[d].kind == DAMACY_AXIS_INDICES) + count += length; + } + if (count > (SIZE_MAX - offset) / sizeof(struct query_index)) + return DAMACY_BUDGET; + char* storage = malloc(offset + (size_t)count * sizeof(struct query_index)); + if (!storage) + return DAMACY_OOM; + memcpy(storage, sample->uri, uri_bytes); + *bounds = (struct damacy_aabb){ .rank = sample->rank }; + memset(axes, 0, DAMACY_MAX_RANK * sizeof(*axes)); + struct query_index* cursor = (void*)(storage + offset); + for (uint8_t d = 0; d < sample->rank; ++d) { + const struct damacy_axis_selection* axis = &sample->axes[d]; + if (axis->kind == DAMACY_AXIS_INTERVAL) { + bounds->dims[d] = axis->interval; + continue; + } + uint32_t n = axis->indices.count; + for (uint32_t i = 0; i < n; ++i) { + int64_t value = axis->indices.values[i]; + if (value < 0 || value == INT64_MAX) { + free(storage); + return DAMACY_INVAL; + } + cursor[i] = (struct query_index){ .source = value, .output = i }; + } + qsort(cursor, n, sizeof(*cursor), compare_indices); + bounds->dims[d] = + (struct damacy_interval){ cursor[0].source, cursor[n - 1].source + 1 }; + axes[d] = (struct query_axis){ .indices = cursor, .count = n }; + cursor += n; + } + *uri = storage; + return DAMACY_OK; +} + +uint32_t +query_lower_bound(const struct query_axis* axis, uint64_t source) +{ + uint32_t lo = 0; + uint32_t hi = axis->count; + while (lo < hi) { + uint32_t mid = lo + (hi - lo) / 2; + if ((uint64_t)axis->indices[mid].source < source) + lo = mid + 1; + else + hi = mid; + } + return lo; +} + +struct selection_span +query_chunk_span(const struct query_axis* axis, + struct damacy_interval bounds, + uint64_t origin, + uint64_t extent) +{ + uint64_t end = origin > UINT64_MAX - extent ? UINT64_MAX : origin + extent; + if (axis->count) { + uint32_t begin = query_lower_bound(axis, origin); + return (struct selection_span){ begin, + query_lower_bound(axis, end) - begin }; + } + uint64_t begin = + origin > (uint64_t)bounds.beg ? origin : (uint64_t)bounds.beg; + if (end > (uint64_t)bounds.end) + end = (uint64_t)bounds.end; + return (struct selection_span){ begin, end > begin ? end - begin : 0 }; +} + +static uint64_t +cell_start_index(uint64_t cell, uint64_t cell_size_index) +{ + return cell > UINT64_MAX / cell_size_index ? UINT64_MAX + : cell * cell_size_index; +} + +enum damacy_status +selection_grid_init(struct selection_grid* grid, + const struct damacy_aabb* bounds_index, + const struct query_axis* axes, + const uint64_t* cell_shape_index, + const uint64_t* clip_beg_cell, + const uint64_t* clip_end_cell, + uint64_t max_cells) +{ + *grid = (struct selection_grid){ .rank = bounds_index->rank, .count = 1 }; + if (!bounds_index->rank || bounds_index->rank > DAMACY_MAX_RANK) + return DAMACY_RANK; + for (uint8_t d = 0; d < bounds_index->rank; ++d) { + if (bounds_index->dims[d].beg < 0 || + bounds_index->dims[d].end <= bounds_index->dims[d].beg || + !cell_shape_index[d]) + return DAMACY_INVAL; + struct selection_grid_axis* axis = &grid->axes[d]; + axis->cell_size_index = cell_shape_index[d]; + axis->span_cell.beg = + (uint64_t)bounds_index->dims[d].beg / cell_shape_index[d]; + axis->span_cell.end = + ((uint64_t)bounds_index->dims[d].end - 1) / cell_shape_index[d] + 1; + if (clip_beg_cell && axis->span_cell.beg < clip_beg_cell[d]) + axis->span_cell.beg = clip_beg_cell[d]; + if (clip_end_cell && axis->span_cell.end > clip_end_cell[d]) + axis->span_cell.end = clip_end_cell[d]; + uint64_t cell_count = axis->span_cell.end > axis->span_cell.beg + ? axis->span_cell.end - axis->span_cell.beg + : 0; + axis->current_cell = axis->span_cell.beg; + if (axes[d].count && cell_count) { + axis->indices = axes[d].indices; + axis->span_entry.beg = query_lower_bound( + &axes[d], cell_start_index(axis->span_cell.beg, cell_shape_index[d])); + axis->span_entry.end = query_lower_bound( + &axes[d], cell_start_index(axis->span_cell.end, cell_shape_index[d])); + axis->current_entry = axis->span_entry.beg; + cell_count = 0; + uint64_t previous_cell = UINT64_MAX; + for (uint32_t entry = axis->span_entry.beg; entry < axis->span_entry.end; + ++entry) { + uint64_t cell = + (uint64_t)axis->indices[entry].source / cell_shape_index[d]; + cell_count += cell != previous_cell; + previous_cell = cell; + } + if (cell_count) + axis->current_cell = + (uint64_t)axis->indices[axis->span_entry.beg].source / + cell_shape_index[d]; + } + if (!cell_count) { + grid->count = 0; + grid->finished = 1; + } else { + if (grid->count > max_cells / cell_count) + return DAMACY_BUDGET; + grid->count *= cell_count; + } + } + return DAMACY_OK; +} + +int +selection_grid_next(struct selection_grid* grid, uint64_t* coordinate_cell) +{ + if (grid->finished) + return 0; + for (uint8_t d = 0; d < grid->rank; ++d) + coordinate_cell[d] = grid->axes[d].current_cell; + for (int d = grid->rank - 1; d >= 0; --d) { + struct selection_grid_axis* axis = &grid->axes[d]; + if (axis->indices) { + while (++axis->current_entry < axis->span_entry.end) { + uint64_t cell = (uint64_t)axis->indices[axis->current_entry].source / + axis->cell_size_index; + if (cell != axis->current_cell) { + axis->current_cell = cell; + return 1; + } + } + axis->current_entry = axis->span_entry.beg; + axis->current_cell = + (uint64_t)axis->indices[axis->span_entry.beg].source / + axis->cell_size_index; + } else { + if (++axis->current_cell < axis->span_cell.end) + return 1; + axis->current_cell = axis->span_cell.beg; + } + } + grid->finished = 1; + return 1; +} diff --git a/src/query/selection.h b/src/query/selection.h new file mode 100644 index 00000000..861b50c9 --- /dev/null +++ b/src/query/selection.h @@ -0,0 +1,82 @@ +#pragma once + +#include "damacy_pipeline.h" + +struct query_index +{ + int64_t source; + uint32_t output; +}; + +struct query_axis +{ + const struct query_index* indices; + uint32_t count; +}; + +struct selection_grid_axis +{ + const struct query_index* indices; + uint64_t cell_size_index; + struct + { + uint64_t beg; + uint64_t end; + } span_cell; + uint64_t current_cell; + struct + { + uint32_t beg; + uint32_t end; + } span_entry; + uint32_t current_entry; +}; + +struct selection_grid +{ + struct selection_grid_axis axes[DAMACY_MAX_RANK]; + uint64_t count; + uint8_t rank; + uint8_t finished; +}; + +struct selection_span +{ + uint64_t begin; + uint64_t count; +}; + +enum damacy_status +query_validate(const struct damacy_sample* sample, + const struct damacy_batch_spec* output, + uint64_t max_plan_bytes); + +int +query_has_indices(const struct damacy_sample* sample); + +enum damacy_status +query_copy(const struct damacy_sample* sample, + char** uri, + struct damacy_aabb* bounds, + struct query_axis* axes); + +uint32_t +query_lower_bound(const struct query_axis* axis, uint64_t source); + +struct selection_span +query_chunk_span(const struct query_axis* axis, + struct damacy_interval bounds, + uint64_t origin, + uint64_t extent); + +enum damacy_status +selection_grid_init(struct selection_grid* grid, + const struct damacy_aabb* bounds_index, + const struct query_axis* axes, + const uint64_t* cell_shape_index, + const uint64_t* clip_beg_cell, + const uint64_t* clip_end_cell, + uint64_t max_cells); + +int +selection_grid_next(struct selection_grid* grid, uint64_t* coordinate_cell); diff --git a/src/render_job/render_job.c b/src/render_job/render_job.c index ac8d57b1..9f95d8f9 100644 --- a/src/render_job/render_job.c +++ b/src/render_job/render_job.c @@ -10,7 +10,10 @@ #include int -render_job_init(struct render_job* job, uint32_t samples_per_batch_cap) +render_job_init(struct render_job* job, + uint32_t samples_per_batch_cap, + uint32_t indices_cap, + uint32_t gather_dims_cap) { memset(job, 0, sizeof(*job)); job->read_ops = (struct read_op*)calloc(DAMACY_MAX_CHUNKS_PER_BATCH, @@ -34,6 +37,25 @@ render_job_init(struct render_job* job, uint32_t samples_per_batch_cap) CUDA_SUCCESS) goto Error; job->d_sample_plans = (void*)(uintptr_t)dptr; + if (indices_cap) { + size_t index_bytes = (size_t)indices_cap * sizeof(*job->indices); + if (cuMemHostAlloc((void**)&job->indices, index_bytes, 0) != CUDA_SUCCESS) + goto Error; + if (cuMemAlloc(&dptr, index_bytes) != CUDA_SUCCESS) + goto Error; + job->d_indices = (void*)(uintptr_t)dptr; + job->indices_cap = indices_cap; + } + if (gather_dims_cap) { + size_t gather_bytes = (size_t)gather_dims_cap * sizeof(*job->gather_dims); + if (cuMemHostAlloc((void**)&job->gather_dims, gather_bytes, 0) != + CUDA_SUCCESS) + goto Error; + if (cuMemAlloc(&dptr, gather_bytes) != CUDA_SUCCESS) + goto Error; + job->d_gather_dims = (void*)(uintptr_t)dptr; + job->gather_dims_cap = gather_dims_cap; + } return 0; Error: render_job_destroy(job, 0); @@ -50,6 +72,14 @@ render_job_destroy(struct render_job* job, int cuda_skip) free(job->read_op_groups); if (!cuda_skip && job->sample_plans) cuMemFreeHost(job->sample_plans); + if (!cuda_skip && job->indices) + cuMemFreeHost(job->indices); + if (!cuda_skip && job->d_indices) + cuMemFree(CUDPTR(job->d_indices)); + if (!cuda_skip && job->gather_dims) + cuMemFreeHost(job->gather_dims); + if (!cuda_skip && job->d_gather_dims) + cuMemFree(CUDPTR(job->d_gather_dims)); path_intern_free(&job->paths); if (!cuda_skip && job->d_sample_plans) cuMemFree(CUDPTR(job->d_sample_plans)); @@ -106,6 +136,8 @@ render_job_reset(struct render_job* job) job->batch_id = 0; job->n_read_op_groups = 0; job->n_sample_plans = 0; + job->n_indices = 0; + job->n_gather_dims = 0; job->n_chunks = 0; job->n_chunks_to_load = 0; job->n_loads_issued = 0; @@ -125,6 +157,10 @@ render_job_dispatch_output(struct render_job* job, uint32_t samples_per_batch) .sample_plans_cap = samples_per_batch, .read_op_groups = job->read_op_groups, .read_op_groups_cap = DAMACY_MAX_CHUNKS_PER_BATCH, + .indices = job->indices, + .indices_cap = job->indices_cap, + .gather_dims = job->gather_dims, + .gather_dims_cap = job->gather_dims_cap, .paths = &job->paths, }; } @@ -134,6 +170,22 @@ render_job_upload_sample_plans(struct render_job* job, void* stream) { if (job->n_sample_plans == 0) return DAMACY_OK; + if (job->n_indices && + cuMemcpyHtoDAsync(CUDPTR(job->d_indices), + job->indices, + (size_t)job->n_indices * sizeof(*job->indices), + (CUstream)stream) != CUDA_SUCCESS) + return DAMACY_CUDA; + if (job->n_gather_dims && + cuMemcpyHtoDAsync(CUDPTR(job->d_gather_dims), + job->gather_dims, + (size_t)job->n_gather_dims * sizeof(*job->gather_dims), + (CUstream)stream) != CUDA_SUCCESS) + return DAMACY_CUDA; + for (uint32_t i = 0; i < job->n_sample_plans; ++i) { + job->sample_plans[i].indices = job->d_indices; + job->sample_plans[i].gather_dims = job->d_gather_dims; + } return cuMemcpyHtoDAsync(CUDPTR(job->d_sample_plans), job->sample_plans, (size_t)job->n_sample_plans * @@ -155,6 +207,8 @@ render_job_commit_plan(struct render_job* job, job->n_chunks_to_load = out->n_chunks_to_load; job->n_loads_issued = out->n_loads_issued; job->n_sample_plans = out->n_sample_plans; + job->n_indices = out->n_indices; + job->n_gather_dims = out->n_gather_dims; job->n_read_op_groups = out->n_read_op_groups; job->n_chunks_dispatched = 0; job->n_groups_dispatched = 0; diff --git a/src/render_job/render_job.h b/src/render_job/render_job.h index 417b0417..ba90a8ea 100644 --- a/src/render_job/render_job.h +++ b/src/render_job/render_job.h @@ -31,6 +31,14 @@ struct render_job struct sample_plan* sample_plans; struct path_intern paths; void* d_sample_plans; + struct gather_index* indices; + void* d_indices; + uint32_t indices_cap; + uint32_t n_indices; + struct gather_dim* gather_dims; + void* d_gather_dims; + uint32_t gather_dims_cap; + uint32_t n_gather_dims; uint32_t n_read_op_groups; uint32_t n_sample_plans; @@ -85,7 +93,10 @@ struct wave_pack_limits }; int -render_job_init(struct render_job* job, uint32_t samples_per_batch_cap); +render_job_init(struct render_job* job, + uint32_t samples_per_batch_cap, + uint32_t indices_cap, + uint32_t gather_dims_cap); void render_job_destroy(struct render_job* job, int cuda_skip); diff --git a/src/wave/wave_budget.c b/src/wave/wave_budget.c index 7d775a47..aa86c870 100644 --- a/src/wave/wave_budget.c +++ b/src/wave/wave_budget.c @@ -132,7 +132,8 @@ gpu_budget_predict(const struct damacy_config* cfg, out->fanout_soa = 2ull * per_wave.fanout_soa; out->nvcomp_temp = nvcomp_temp; out->batch_metadata = - 2ull * (uint64_t)cfg->samples_per_batch * sizeof(struct sample_plan); + 2ull * ((uint64_t)cfg->samples_per_batch * sizeof(struct sample_plan) + + dispatch_index_storage_bytes(cfg)); out->total = out->dev_compressed + out->dev_decompressed + out->blosc1_meta + out->fanout_soa + out->nvcomp_temp + out->batch_metadata; return DAMACY_OK; diff --git a/src/wave/wave_pool.c b/src/wave/wave_pool.c index 5daf22ec..3b823bd1 100644 --- a/src/wave/wave_pool.c +++ b/src/wave/wave_pool.c @@ -651,16 +651,22 @@ set_assemble_chunk_base(struct assemble_chunk* out, out->src_base_byte_off = (uint64_t)c->dev_decompressed_offset; out->sample_idx_in_batch = c->sample_idx_in_batch; out->is_fill = c->is_fill | force_fill; + out->gather_offset = c->gather_offset; memcpy(out->chunk_d, c->chunk_d, spatial_rank * sizeof(*c->chunk_d)); } static uint64_t chunk_output_elements(const struct chunk_plan* c, const struct sample_plan* sp, + const struct gather_dim* gather_dims, uint8_t spatial_rank) { uint64_t out = 1; for (uint8_t d = 0; d < spatial_rank; ++d) { + if (sp->indexed) { + out *= gather_dims[c->gather_offset + d].count; + continue; + } int64_t chunk_shape = (int64_t)sp->dims[d].chunk_shape; int64_t origin = (int64_t)c->chunk_d[d] * chunk_shape - sp->dims[d].aabb_lo_relative; @@ -687,12 +693,15 @@ build_assemble_meta(const struct wave_pool* wp, struct damacy_wave* wave) set_assemble_chunk_base( &wave->h_assemble_chunks[i], c, spatial_rank, wp->bypass_decode); - uint32_t bpc = assemble_blocks_per_chunk(spatial_rank, sp->dims); + uint64_t elements = + chunk_output_elements(c, sp, job->gather_dims, spatial_rank); + uint32_t bpc = sp->indexed + ? assemble_gather_blocks(elements) + : assemble_blocks_per_chunk(spatial_rank, sp->dims); if (bpc > max_bpc) max_bpc = bpc; - wave->assemble_out_bytes += - chunk_output_elements(c, sp, spatial_rank) * (uint64_t)bpe; + wave->assemble_out_bytes += elements * (uint64_t)bpe; } if (max_bpc == 0) max_bpc = 1; @@ -706,7 +715,8 @@ decode_anchor_reserve(struct wave_pool* wp, { size_t prev_idx = wp->decode_done_ring_idx; *anchor_idx = (prev_idx + 1) % countof(wp->decode_done_ring); - wave->prev_decode_anchor = wp->decode_done_ring[prev_idx]; + wave->prev_decode_anchor = + wp->decode_done_recorded ? wp->decode_done_ring[prev_idx] : NULL; return wp->decode_done_ring[*anchor_idx]; } @@ -714,6 +724,7 @@ static void decode_anchor_commit(struct wave_pool* wp, size_t anchor_idx) { wp->decode_done_ring_idx = (uint8_t)anchor_idx; + wp->decode_done_recorded = 1; } static enum damacy_status diff --git a/src/wave/wave_pool.h b/src/wave/wave_pool.h index bd0fa2e1..384830af 100644 --- a/src/wave/wave_pool.h +++ b/src/wave/wave_pool.h @@ -34,6 +34,7 @@ struct wave_pool CUevent decode_done_ring[4]; uint8_t decode_done_ring_idx; + uint8_t decode_done_recorded; struct decoder_zstd* zstd_decoder; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c2394f3f..34f54287 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -52,6 +52,7 @@ add_damacy_test( 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 dispatch_utils log) +add_damacy_test(test_selection query_selection 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) diff --git a/tests/test_assemble.c b/tests/test_assemble.c index 01ff2827..3b668f31 100644 --- a/tests/test_assemble.c +++ b/tests/test_assemble.c @@ -243,7 +243,6 @@ run_scenario(int rank, sp.rank = (uint8_t)rank; sp.src_dtype = (uint8_t)src_dtype; sp.sample_dst_off_elems = 0; - sp.chunk_count = n_chunks; for (int d = 0; d < rank; ++d) { sp.dims[d].chunk_shape = S[d]; sp.dims[d].chunk_grid_extent = N[d]; @@ -555,7 +554,6 @@ run_fill_scenario(enum dtype src_dtype, struct sample_plan sp = { 0 }; sp.rank = (uint8_t)rank; sp.src_dtype = (uint8_t)src_dtype; - sp.chunk_count = n_chunks; for (int d = 0; d < rank; ++d) { sp.dims[d].chunk_shape = S[d]; sp.dims[d].chunk_grid_extent = N[d]; @@ -747,7 +745,6 @@ test_rank3_all_fill(void) struct sample_plan sp = { 0 }; sp.rank = (uint8_t)rank; sp.src_dtype = (uint8_t)src_dtype; - sp.chunk_count = n_chunks; for (int d = 0; d < rank; ++d) { sp.dims[d].chunk_shape = S[d]; sp.dims[d].chunk_grid_extent = N[d]; diff --git a/tests/test_cpu_pipeline.c b/tests/test_cpu_pipeline.c index d064b6f0..72e83ece 100644 --- a/tests/test_cpu_pipeline.c +++ b/tests/test_cpu_pipeline.c @@ -88,7 +88,9 @@ 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 } } } + .rank = 2, + .axes = { { .kind = DAMACY_AXIS_INTERVAL, .interval = { y, y + rows } }, + { .kind = DAMACY_AXIS_INTERVAL, .interval = { x, x + cols } } } }; } @@ -433,9 +435,70 @@ test_bfloat_rounding_and_fill(void) return 0; } +static int +test_invalid_axes_rejected(void) +{ + struct components c = { 0 }; + EXPECT(create_components(&c) == 0); + EXPECT(start_pipeline(&c, DAMACY_F32, 2, 2, 1) == 0); + struct damacy_sample query = sample("unused", 0, 0, 2, 2); + int kinds[] = { 0, -1, 3 }; + for (size_t i = 0; i < sizeof(kinds) / sizeof(*kinds); ++i) { + query.axes[1].kind = (enum damacy_axis_kind)kinds[i]; + struct damacy_push_result pushed = damacy_push( + c.pipeline, (struct damacy_sample_slice){ &query, &query + 1 }); + EXPECT(pushed.status == DAMACY_INVAL); + EXPECT(pushed.unconsumed.beg == &query && + pushed.unconsumed.end == &query + 1); + } + destroy_components(&c); + return 0; +} + +static int +test_index_arrays_copied(void) +{ + char root[] = "/tmp/damacy_indices_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, "float32", 0, "none") == 0); + struct components c = { 0 }; + EXPECT(create_components(&c) == 0); + EXPECT(start_pipeline(&c, DAMACY_F32, 3, 4, 1) == 0); + int64_t rows[] = { 4, 0, 4 }, cols[] = { 10, 0, 3, 10 }; + struct damacy_sample query = { + .uri = uri, + .rank = 2, + .axes = { { .kind = DAMACY_AXIS_INDICES, .indices = { rows, 3 } }, + { .kind = DAMACY_AXIS_INDICES, .indices = { cols, 4 } } } + }; + struct damacy_push_result pushed = + damacy_push(c.pipeline, (struct damacy_sample_slice){ &query, &query + 1 }); + EXPECT(pushed.status == DAMACY_OK && pushed.unconsumed.beg == &query + 1); + memset(rows, 0, sizeof(rows)); + memset(cols, 0, sizeof(cols)); + memset(uri, 0, sizeof(uri)); + struct damacy_batch* batch = NULL; + EXPECT(damacy_pop(c.pipeline, &batch) == DAMACY_OK); + destroy_components(&c); + struct damacy_batch_info info; + damacy_batch_info(batch, &info); + const float expected[] = { 54, 44, 47, 54, 10, 0, 3, 10, 54, 44, 47, 54 }; + EXPECT(info.shape[0] == 1 && info.shape[1] == 3 && info.shape[2] == 4); + EXPECT(!memcmp(info.data, expected, sizeof(expected))); + damacy_batch_release(batch); + fixture_rm_tree(root); + return 0; +} + int main(void) { + RUN(test_invalid_axes_rejected); + RUN(test_index_arrays_copied); RUN(test_codecs_and_types); RUN(test_owned_plan); RUN(test_shared_metadata); diff --git a/tests/test_damacy.c b/tests/test_damacy.c index e8b5580a..b479f600 100644 --- a/tests/test_damacy.c +++ b/tests/test_damacy.c @@ -75,9 +75,13 @@ mk_cfg(const char* root, uint32_t samples_per_batch, int64_t sy, int64_t sx) static struct damacy_sample mk_sample(const char* uri, int64_t y0, int64_t y1, int64_t x0, int64_t x1) { - struct damacy_sample s = { .uri = uri, .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = y0, .end = y1 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = x0, .end = x1 }; + struct damacy_sample s = { .uri = uri, .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = y0, .end = y1 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = x0, .end = x1 } }; return s; } diff --git a/tests/test_damacy_blosc.c b/tests/test_damacy_blosc.c index efc88971..bbef48c9 100644 --- a/tests/test_damacy_blosc.c +++ b/tests/test_damacy_blosc.c @@ -78,9 +78,13 @@ mk_cfg(const char* root, uint32_t samples_per_batch, int64_t sy, int64_t sx) static struct damacy_sample mk_sample(const char* uri, int64_t y0, int64_t y1, int64_t x0, int64_t x1) { - struct damacy_sample s = { .uri = uri, .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = y0, .end = y1 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = x0, .end = x1 }; + struct damacy_sample s = { .uri = uri, .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = y0, .end = y1 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = x0, .end = x1 } }; return s; } diff --git a/tests/test_damacy_caps.c b/tests/test_damacy_caps.c index 6c72d4e0..5ec60c88 100644 --- a/tests/test_damacy_caps.c +++ b/tests/test_damacy_caps.c @@ -31,6 +31,10 @@ // test_sample_shape_mismatch_rejected // — push a sample whose aabb extent != // cfg.sample_shape; expect INVAL +// test_index_capacity — too small a nonzero max_index_bytes +// fails create (BUDGET); with 0, push +// rejects an indexed sample (BUDGET) +// and the pipeline keeps running // test_resolver_minimum_one_chunk — budget barely fitting a // single-chunk wave produces a valid // instance and surfaces tight geometry @@ -84,9 +88,13 @@ mk_cfg(const char* root, uint32_t samples_per_batch, int64_t sy, int64_t sx) static struct damacy_sample mk_sample(const char* uri, int64_t y0, int64_t y1, int64_t x0, int64_t x1) { - struct damacy_sample s = { .uri = uri, .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = y0, .end = y1 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = x0, .end = x1 }; + struct damacy_sample s = { .uri = uri, .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = y0, .end = y1 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = x0, .end = x1 } }; return s; } @@ -266,14 +274,22 @@ test_pool_reserve_fits_default_budget(void) struct damacy* d = NULL; EXPECT(damacy_create(&cfg, &d) == DAMACY_OK); - struct damacy_sample samples[20]; + struct damacy_sample samples[20] = { 0 }; for (int i = 0; i < 20; ++i) { samples[i].uri = p; - samples[i].aabb.rank = 4; - samples[i].aabb.dims[0] = (struct damacy_interval){ 0, 16 }; - samples[i].aabb.dims[1] = (struct damacy_interval){ 0, 1 }; - samples[i].aabb.dims[2] = (struct damacy_interval){ 0, 256 }; - samples[i].aabb.dims[3] = (struct damacy_interval){ 0, 256 }; + samples[i].rank = 4; + samples[i].axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { 0, 16 } }; + samples[i].axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { 0, 1 } }; + samples[i].axes[2] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { 0, 256 } }; + samples[i].axes[3] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { 0, 256 } }; } struct damacy_sample_slice slice = { .beg = samples, .end = samples + 20 }; EXPECT(damacy_push(d, slice).status == DAMACY_OK); @@ -376,6 +392,48 @@ test_sample_shape_mismatch_rejected(void) return 0; } +static int +test_index_capacity(void) +{ + char root[64]; + EXPECT(mkdtemp_root(root, sizeof root) == 0); + char p[256]; + snprintf(p, sizeof p, "%s/foo", root); + int64_t shape[2] = { 8, 16 }, inner[2] = { 4, 8 }, shard[2] = { 8, 16 }; + EXPECT(fixture_write_zarr_codec( + p, shape, inner, shard, 2, "uint16", 0, "blosc-zstd") == 0); + + // One sample of shape (8, 16) holds at most 8 + 16 indices. + struct damacy_config cfg = mk_cfg(root, 1, 8, 16); + cfg.tuning.max_index_bytes = 8 * (8 + 16) - 1; + struct damacy* d = NULL; + EXPECT(damacy_create(&cfg, &d) == DAMACY_BUDGET); + EXPECT(d == NULL); + + cfg.tuning.max_index_bytes = 0; + EXPECT(damacy_create(&cfg, &d) == DAMACY_OK); + int64_t rows[8] = { 7, 6, 5, 4, 3, 2, 1, 0 }; + struct damacy_sample s[2] = { mk_sample(p, 0, 8, 0, 16), + mk_sample(p, 0, 8, 0, 16) }; + s[1].axes[0] = (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INDICES, + .indices = { rows, 8 } }; + struct damacy_push_result pr = + damacy_push(d, (struct damacy_sample_slice){ s, s + 2 }); + EXPECT(pr.status == DAMACY_BUDGET); + EXPECT(pr.unconsumed.beg == s + 1 && pr.unconsumed.end == s + 2); + struct damacy_batch* b = NULL; + EXPECT(damacy_pop(d, &b) == DAMACY_OK); + damacy_release(d, b); + pr = damacy_push(d, (struct damacy_sample_slice){ s, s + 1 }); + EXPECT(pr.status == DAMACY_OK); + EXPECT(damacy_pop(d, &b) == DAMACY_OK); + damacy_release(d, b); + + damacy_destroy(d); + fixture_rm_tree(root); + return 0; +} + // Resolver picks per-wave near the minimum when the budget barely // fits. Reports gpu_bytes_committed back through stats so users can // observe the resolved size. @@ -556,6 +614,7 @@ main(void) RUN(test_cache_floors_validated_at_create); RUN(test_oversized_sample_shard_count_rejected); RUN(test_sample_shape_mismatch_rejected); + RUN(test_index_capacity); RUN(test_resolver_minimum_one_chunk); log_info("all tests passed"); return 0; diff --git a/tests/test_damacy_ctx.c b/tests/test_damacy_ctx.c index 5e576a95..16988989 100644 --- a/tests/test_damacy_ctx.c +++ b/tests/test_damacy_ctx.c @@ -211,9 +211,13 @@ test_explicit_device_does_not_leak_ctx_between_calls(void) EXPECT(cuCtxGetCurrent(&cur) == CUDA_SUCCESS); EXPECT(cur == caller); - struct damacy_sample s = { .uri = p, .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 4 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 8 }; + struct damacy_sample s = { .uri = p, .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 4 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 8 } }; struct damacy_sample_slice slice = { .beg = &s, .end = &s + 1 }; struct damacy_push_result pr = damacy_push(d, slice); EXPECT(pr.status == DAMACY_OK); @@ -272,9 +276,13 @@ test_pop_error_path_restores_caller_ctx(void) // Push succeeds (shape matches cfg), pop's planner rejects the AABB // with DAMACY_INVAL after ctx_guard_enter. - struct damacy_sample oob = { .uri = p, .aabb = { .rank = 2 } }; - oob.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 4 }; - oob.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 9999 }; + struct damacy_sample oob = { .uri = p, .rank = 2 }; + oob.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 4 } }; + oob.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 9999 } }; struct damacy_sample_slice slice = { .beg = &oob, .end = &oob + 1 }; struct damacy_push_result pr = damacy_push(d, slice); EXPECT(pr.status == DAMACY_OK); diff --git a/tests/test_lookahead.c b/tests/test_lookahead.c index 726d53ce..989cd1ad 100644 --- a/tests/test_lookahead.c +++ b/tests/test_lookahead.c @@ -11,9 +11,9 @@ mk_sample(const char* uri, int64_t lo, int64_t hi) { struct damacy_sample s = { 0 }; s.uri = uri; - s.aabb.rank = 1; - s.aabb.dims[0].beg = lo; - s.aabb.dims[0].end = hi; + s.rank = 1; + s.axes[0] = (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { lo, hi } }; return s; } diff --git a/tests/test_planner.c b/tests/test_planner.c index 5617b58b..e436eb23 100644 --- a/tests/test_planner.c +++ b/tests/test_planner.c @@ -414,7 +414,6 @@ test_single_chunk_aligned(void) EXPECT(sp->rank == 2); EXPECT(sp->batch_pool_slot == 0); EXPECT(sp->sample_idx_in_batch == 0); - EXPECT(sp->chunk_count == 1); EXPECT(sp->dims[0].chunk_shape == 2 && sp->dims[1].chunk_shape == 4); EXPECT(sp->dims[0].chunk_grid_extent == 1 && sp->dims[1].chunk_grid_extent == 1); @@ -489,7 +488,6 @@ test_multi_chunk_partial(void) struct sample_plan* sp = &samples[0]; EXPECT(sp->rank == 2); - EXPECT(sp->chunk_count == 4); EXPECT(sp->dims[0].chunk_grid_extent == 2 && sp->dims[1].chunk_grid_extent == 2); EXPECT(sp->dims[0].aabb_lo_relative == 1 && @@ -542,7 +540,6 @@ test_two_samples_indices(void) EXPECT(samples[0].sample_dst_off_elems == 0); EXPECT(samples[1].sample_dst_off_elems == 8); // sample_idx=1 * stride[0]=8 - EXPECT(samples[0].chunk_count == 1 && samples[1].chunk_count == 1); fixture_destroy(&f); return 0; @@ -1480,9 +1477,109 @@ test_coalesce_cross_sample(void) return 0; } +static int +test_owned_index_plan(void) +{ + const uint64_t offsets[4] = { 0, 100, 200, 300 }; + const uint64_t nbytes[4] = { 32, 32, 32, 32 }; + struct fixture f = { 0 }; + EXPECT(fixture_init(&f, offsets, nbytes) == 0); + struct planner_sample samples[2]; + EXPECT(mk_sample(&f, 0, 4, 0, 8, &samples[0]) == 0); + struct query_index rows[] = { { 0, 1 }, { 3, 0 }, { 3, 2 } }; + struct query_index cols[] = { { 0, 2 }, { 3, 1 }, { 7, 0 }, { 7, 3 } }; + samples[0].axes[0] = (struct query_axis){ rows, 3 }; + samples[0].axes[1] = (struct query_axis){ cols, 4 }; + samples[1] = samples[0]; + struct damacy_batch_spec output = { .dtype = DAMACY_F32, + .sample_rank = 2, + .sample_shape = { 3, 4 }, + .samples_per_batch = 2 }; + struct damacy_plan_limits limits = { .max_chunks = 8, + .max_chunk_bytes = 1024, + .max_shards_per_sample = 1, + .max_plan_bytes = 1 << 20 }; + struct prepared_plan* plan = NULL; + EXPECT(prepared_plan_build(f.array_meta_cache, + f.shard_index_cache, + samples, + 2, + &output, + &limits, + &plan) == DAMACY_OK); + EXPECT(plan->n_chunks == 4 && plan->n_uses == 8 && plan->n_arrays == 1); + EXPECT(plan->regions[0].operation == PLAN_GATHER); + EXPECT(plan->regions[0].axes[0].indices != rows); + EXPECT(plan->regions[0].axes[1].indices != cols); + limits.max_plan_bytes = plan->allocated_bytes - 1; + struct prepared_plan* rejected = NULL; + EXPECT(prepared_plan_build(f.array_meta_cache, + f.shard_index_cache, + samples, + 2, + &output, + &limits, + &rejected) == DAMACY_BUDGET); + EXPECT(!rejected); + ++limits.max_plan_bytes; + EXPECT(prepared_plan_build(f.array_meta_cache, + f.shard_index_cache, + samples, + 2, + &output, + &limits, + &rejected) == DAMACY_OK); + prepared_plan_destroy(rejected); + memset(rows, 0, sizeof(rows)); + memset(cols, 0, sizeof(cols)); + fixture_destroy(&f); + EXPECT(plan->regions[0].axes[0].indices[1].source == 3); + EXPECT(plan->regions[1].axes[1].indices[2].output == 0); + EXPECT(plan->arrays[0].metadata.inner_chunk_shape[0] == 2); + EXPECT(!strcmp(plan->chunks[0].path, "foo/c/0/0")); + struct dispatch_output dispatch = { + .read_ops = (struct read_op[8]){ 0 }, + .read_ops_cap = 8, + .chunk_plans = (struct chunk_plan[8]){ 0 }, + .chunk_plans_cap = 8, + .sample_plans = (struct sample_plan[2]){ 0 }, + .sample_plans_cap = 2, + .read_op_groups = (struct read_op_group[8]){ 0 }, + .read_op_groups_cap = 8, + .indices = (struct gather_index[14]){ 0 }, + .indices_cap = 13, + .gather_dims = (struct gather_dim[16]){ 0 }, + .gather_dims_cap = 15, + .paths = &(struct path_intern){ 0 } + }; + struct dispatch_scratch scratch = { 0 }; + EXPECT(dispatch_plan_build(plan, 0, PAGE, PAGE, 8, &dispatch, &scratch) == + DAMACY_BUDGET); + dispatch.indices_cap = 14; + EXPECT(dispatch_plan_build(plan, 0, PAGE, PAGE, 8, &dispatch, &scratch) == + DAMACY_BUDGET); + dispatch.gather_dims_cap = 16; + EXPECT(dispatch_plan_build(plan, 0, PAGE, PAGE, 8, &dispatch, &scratch) == + DAMACY_OK); + EXPECT(dispatch.n_indices == 14 && dispatch.n_gather_dims == 16); + EXPECT(dispatch.sample_plans[0].indexed); + uint64_t elements = 0; + for (uint32_t i = 0; i < dispatch.n_chunk_plans; ++i) { + const struct gather_dim* gather = + dispatch.gather_dims + dispatch.chunk_plans[i].gather_offset; + elements += (uint64_t)gather[0].count * gather[1].count; + } + EXPECT(elements == 24); + path_intern_free(dispatch.paths); + dispatch_scratch_destroy(&scratch); + prepared_plan_destroy(plan); + return 0; +} + int main(void) { + RUN(test_owned_index_plan); RUN(test_single_chunk_aligned); RUN(test_multi_chunk_partial); RUN(test_two_samples_indices); diff --git a/tests/test_prefetcher.c b/tests/test_prefetcher.c index 0963f119..6489f6ad 100644 --- a/tests/test_prefetcher.c +++ b/tests/test_prefetcher.c @@ -157,9 +157,13 @@ test_single_sample_walks_all_stages(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -189,9 +193,13 @@ test_dedup_across_samples(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; for (int i = 0; i < 4; ++i) EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, (uint64_t)i) == 0); @@ -220,9 +228,13 @@ test_error_propagates(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "missing", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "missing", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -259,9 +271,13 @@ test_multi_shard_sample_warms_all_shards(void) int64_t shard[2] = { 16, 32 }; EXPECT(fixture_setup_layout(&fx, "blosc-zstd", shape, inner, shard) == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 32 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 64 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 64 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -290,9 +306,13 @@ test_pop_ready_returns_completed_sample(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -320,9 +340,13 @@ test_pop_ready_surfaces_error_state(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "missing", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "missing", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -344,9 +368,13 @@ test_pop_ready_recycles_slot(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; for (int round = 0; round < 3; ++round) { EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, (uint64_t)round) == @@ -385,9 +413,13 @@ test_sample_gate_ready_after_drain(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -407,9 +439,13 @@ test_sample_gate_error_on_failed_sample(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "missing", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "missing", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -428,9 +464,13 @@ test_distinct_samples_get_separate_gates(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 1) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -451,9 +491,13 @@ test_pop_drops_sample_gate(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -474,9 +518,13 @@ test_sample_gate_lives_until_pop(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -500,9 +548,13 @@ test_admit_fail_releases_batch_entry(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; for (int i = 0; i < 8; ++i) { char uri[16]; @@ -555,9 +607,13 @@ test_take_wave_yields_ticket(void) EXPECT(prefetcher_in_flight(fx.p) == 0); EXPECT(prefetcher_ready_prefix_count(fx.p) == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 1) == 0); @@ -597,9 +653,13 @@ test_take_wave_caps_to_ready_prefix(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -624,9 +684,13 @@ test_ready_wave_waits_for_prefix(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 1) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -685,9 +749,13 @@ test_owner_capacity_saturation_surfaces_error(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; for (uint64_t i = 0; i < 8; ++i) EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, i) == 0); @@ -722,9 +790,13 @@ test_owner_table_recycles_after_pop(void) struct fixture fx = { 0 }; EXPECT(fixture_setup(&fx, "blosc-zstd") == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; for (uint64_t i = 0; i < 8; ++i) EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, i) == 0); @@ -766,9 +838,13 @@ test_missing_shard_reaches_ready(void) snprintf(shard_path, sizeof shard_path, "%s/foo/c/0/0", fx.root); EXPECT(unlink(shard_path) == 0); - struct damacy_sample s = { .uri = "foo", .aabb = { .rank = 2 } }; - s.aabb.dims[0] = (struct damacy_interval){ .beg = 0, .end = 16 }; - s.aabb.dims[1] = (struct damacy_interval){ .beg = 0, .end = 32 }; + struct damacy_sample s = { .uri = "foo", .rank = 2 }; + s.axes[0] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 16 } }; + s.axes[1] = + (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { .beg = 0, .end = 32 } }; EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &s, 0) == 0); EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); @@ -788,9 +864,48 @@ test_missing_shard_reaches_ready(void) return 0; } +static int +test_sparse_indices_survive_prefetch(void) +{ + struct fixture fx = { 0 }; + int64_t shape[] = { 32, 32 }, chunks[] = { 4, 4 }, shards[] = { 8, 8 }; + EXPECT(fixture_setup_layout(&fx, "none", shape, chunks, shards) == 0); + int64_t rows[] = { 31, 0, 31 }, cols[] = { 31, 0 }; + struct damacy_sample sample = { + .uri = "foo", + .rank = 2, + .axes = { { .kind = DAMACY_AXIS_INDICES, .indices = { rows, 3 } }, + { .kind = DAMACY_AXIS_INDICES, .indices = { cols, 2 } } } + }; + EXPECT(lookahead_push_with_sample_seq(&fx.lookahead, &sample, 0) == 0); + memset(rows, 0, sizeof(rows)); + memset(cols, 0, sizeof(cols)); + EXPECT(prefetcher_drain(fx.p) == DAMACY_OK); + struct prefetch_cache_stats stats; + prefetch_cache_stats_get(fx.shard_index_cache, &stats); + EXPECT(stats.size == 4); + struct prefetcher_ready ready = { 0 }; + EXPECT(prefetcher_pop_ready(fx.p, &ready)); + EXPECT(ready.result == PREFETCHER_RESULT_READY && ready.n_shards == 4); + fixture_teardown(&fx); + EXPECT(ready.aabb.dims[0].beg == 0 && ready.aabb.dims[0].end == 32); + EXPECT(ready.axes[0].count == 3 && ready.axes[1].count == 2); + EXPECT(ready.axes[0].indices[0].source == 0 && + ready.axes[0].indices[0].output == 1); + EXPECT(ready.axes[0].indices[1].source == 31 && + ready.axes[0].indices[1].output == 0); + EXPECT(ready.axes[0].indices[2].source == 31 && + ready.axes[0].indices[2].output == 2); + EXPECT(ready.axes[1].indices[1].source == 31 && + ready.axes[1].indices[1].output == 0); + prefetcher_ready_free(&ready); + return 0; +} + int main(void) { + RUN(test_sparse_indices_survive_prefetch); RUN(test_single_sample_walks_all_stages); RUN(test_dedup_across_samples); RUN(test_error_propagates); diff --git a/tests/test_selection.c b/tests/test_selection.c new file mode 100644 index 00000000..29477355 --- /dev/null +++ b/tests/test_selection.c @@ -0,0 +1,178 @@ +#include "expect.h" +#include "query/selection.h" + +#include +#include + +static int +test_sparse_grid(void) +{ + int64_t rows[] = { INT64_C(1) << 40, 3, 3, 8 }; + int64_t cols[] = { 9, 0, 9 }; + struct damacy_sample sample = { + .uri = "array", + .rank = 2, + .axes = { { .kind = DAMACY_AXIS_INDICES, .indices = { rows, 4 } }, + { .kind = DAMACY_AXIS_INDICES, .indices = { cols, 3 } } } + }; + struct damacy_batch_spec output = { .sample_rank = 2, + .sample_shape = { 4, 3 }, + .samples_per_batch = 1 }; + EXPECT(query_validate(&sample, &output, 7 * sizeof(struct query_index)) == + DAMACY_OK); + EXPECT(query_validate(&sample, &output, 7 * sizeof(struct query_index) - 1) == + DAMACY_BUDGET); + struct damacy_aabb bounds; + struct query_axis axes[DAMACY_MAX_RANK]; + char* storage = NULL; + EXPECT(query_copy(&sample, &storage, &bounds, axes) == DAMACY_OK); + memset(rows, 0, sizeof(rows)); + memset(cols, 0, sizeof(cols)); + EXPECT(!strcmp(storage, "array")); + EXPECT(axes[0].indices[0].source == 3 && axes[0].indices[0].output == 1); + EXPECT(axes[0].indices[1].source == 3 && axes[0].indices[1].output == 2); + EXPECT(axes[0].indices[3].source == (INT64_C(1) << 40)); + uint64_t shape[] = { 4, 4 }; + struct selection_grid grid; + EXPECT(selection_grid_init(&grid, &bounds, axes, shape, NULL, NULL, 5) == + DAMACY_BUDGET); + EXPECT(selection_grid_init(&grid, &bounds, axes, shape, NULL, NULL, 6) == + DAMACY_OK); + EXPECT(grid.count == 6); + uint64_t expected_rows[] = { 0, 2, UINT64_C(1) << 38 }; + uint64_t coordinate[DAMACY_MAX_RANK]; + for (uint32_t i = 0; i < 6; ++i) { + EXPECT(selection_grid_next(&grid, coordinate)); + EXPECT(coordinate[0] == expected_rows[i / 2]); + EXPECT(coordinate[1] == (i % 2 ? 2 : 0)); + } + EXPECT(!selection_grid_next(&grid, coordinate)); + uint64_t begin[] = { 1, 0 }, end[] = { 3, 1 }; + EXPECT(selection_grid_init(&grid, &bounds, axes, shape, begin, end, 1) == + DAMACY_OK); + EXPECT(grid.count == 1); + EXPECT(selection_grid_next(&grid, coordinate)); + EXPECT(coordinate[0] == 2 && coordinate[1] == 0); + EXPECT(!selection_grid_next(&grid, coordinate)); + struct selection_span span = query_chunk_span(&axes[0], bounds.dims[0], 0, 4); + EXPECT(span.begin == 0 && span.count == 2); + span = query_chunk_span(&axes[0], bounds.dims[0], 4, 4); + EXPECT(span.count == 0); + free(storage); + return 0; +} + +static int +test_invalid_indices(void) +{ + int64_t values[] = { 0, 1 }; + struct damacy_sample sample = { .uri = "array", + .rank = 1, + .axes = { { .kind = DAMACY_AXIS_INDICES, + .indices = { values, 2 } } } }; + struct damacy_batch_spec output = { .sample_rank = 1, .sample_shape = { 2 } }; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_OK); + values[0] = -1; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_INVAL); + values[0] = INT64_MAX; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_INVAL); + values[0] = 0; + sample.axes[0].indices.count = 0; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_INVAL); + sample.axes[0].indices = (struct damacy_index_array){ NULL, 2 }; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_INVAL); + sample.axes[0].indices = (struct damacy_index_array){ values, 1 }; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_INVAL); + return 0; +} + +static int +test_invalid_axis_selections(void) +{ + int64_t values[] = { 1, 0 }; + struct damacy_axis_selection invalid[] = { + { .interval = { 3, 5 } }, + { .kind = (enum damacy_axis_kind) - 1, .interval = { 3, 5 } }, + { .kind = (enum damacy_axis_kind)3, .interval = { 3, 5 } }, + { .kind = DAMACY_AXIS_INDICES }, + { .kind = DAMACY_AXIS_INDICES, .indices = { NULL, 2 } }, + { .kind = DAMACY_AXIS_INDICES, .indices = { values, 0 } }, + { .kind = DAMACY_AXIS_INTERVAL, .interval = { 2, 2 } }, + { .kind = DAMACY_AXIS_INTERVAL, .interval = { 3, 2 } }, + { .kind = DAMACY_AXIS_INTERVAL, .interval = { -1, 1 } }, + }; + struct damacy_sample sample = { .uri = "array", .rank = 1 }; + struct damacy_batch_spec output = { .sample_rank = 1, .sample_shape = { 2 } }; + struct damacy_aabb bounds; + struct query_axis axes[DAMACY_MAX_RANK]; + char* storage = NULL; + for (size_t i = 0; i < sizeof(invalid) / sizeof(*invalid); ++i) { + sample.axes[0] = invalid[i]; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_INVAL); + EXPECT(query_copy(&sample, &storage, &bounds, axes) == DAMACY_INVAL); + EXPECT(!storage); + } + sample.axes[0] = (struct damacy_axis_selection){ .kind = DAMACY_AXIS_INTERVAL, + .interval = { 3, 5 } }; + EXPECT(query_validate(&sample, &output, 0) == DAMACY_OK); + EXPECT(query_copy(&sample, &storage, &bounds, axes) == DAMACY_OK); + EXPECT(bounds.rank == 1 && bounds.dims[0].beg == 3 && + bounds.dims[0].end == 5); + EXPECT(!axes[0].indices && !axes[0].count); + free(storage); + sample.rank = 0; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_RANK); + sample.rank = DAMACY_MAX_RANK + 1; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_RANK); + sample.rank = 1; + output.sample_rank = 2; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_RANK); + output.sample_rank = 1; + output.sample_shape[0] = 3; + EXPECT(query_validate(&sample, &output, 1024) == DAMACY_INVAL); + return 0; +} + +static int +test_mixed_grid(void) +{ + int64_t values[] = { 10, 1, 10 }; + struct damacy_sample sample = { + .uri = "array", + .rank = 2, + .axes = { { .kind = DAMACY_AXIS_INDICES, .indices = { values, 3 } }, + { .kind = DAMACY_AXIS_INTERVAL, .interval = { 3, 9 } } } + }; + struct damacy_aabb bounds; + struct query_axis axes[DAMACY_MAX_RANK]; + char* storage = NULL; + EXPECT(query_copy(&sample, &storage, &bounds, axes) == DAMACY_OK); + struct selection_grid grid; + uint64_t shape[] = { 4, 4 }; + EXPECT(selection_grid_init(&grid, &bounds, axes, shape, NULL, NULL, 6) == + DAMACY_OK); + EXPECT(grid.count == 6); + uint64_t coordinate[DAMACY_MAX_RANK]; + for (uint32_t i = 0; i < 6; ++i) { + EXPECT(selection_grid_next(&grid, coordinate)); + EXPECT(coordinate[0] == (i < 3 ? 0 : 2)); + EXPECT(coordinate[1] == i % 3); + } + EXPECT(!selection_grid_next(&grid, coordinate)); + struct selection_span span = query_chunk_span(&axes[1], bounds.dims[1], 0, 4); + EXPECT(span.begin == 3 && span.count == 1); + span = query_chunk_span(&axes[1], bounds.dims[1], 8, 4); + EXPECT(span.begin == 8 && span.count == 1); + free(storage); + return 0; +} + +int +main(void) +{ + RUN(test_sparse_grid); + RUN(test_invalid_indices); + RUN(test_invalid_axis_selections); + RUN(test_mixed_grid); + return 0; +}