From 333c86c5da01d98b7799d4c1a02b1f460503ee9e Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Sun, 4 Oct 2026 10:12:15 -0700 Subject: [PATCH 1/4] Map packed stream input directly into edge chunks --- README.md | 11 +- bench/bench_input.c | 31 +--- bench/bench_input.h | 4 +- bench/bench_report.c | 2 +- bench/bench_util.c | 11 +- bench/test_bench_images.py | 14 +- docs/buffered-writer.md | 9 +- scripts/datasets/README.md | 7 +- scripts/datasets/run.py | 8 +- scripts/datasets/verify_output.py | 13 +- scripts/sweep/microscopy_data.py | 8 +- scripts/sweep/sweep.py | 12 +- scripts/sweep/test_microscopy.py | 18 +++ src/cpu/stream.body.c | 21 ++- src/cpu/stream.c | 7 +- src/cpu/transpose.cpp | 56 ++++++- src/cpu/transpose.h | 3 +- src/gpu/schedule.c | 11 ++ src/gpu/stream.ingest.c | 19 +-- src/gpu/stream.init.c | 19 ++- src/gpu/transpose.cu | 103 ++++++++++++- src/gpu/transpose.h | 13 ++ src/lod/lod_plan.c | 4 +- src/multiarray/stream.c | 7 +- src/stream/config.c | 9 +- src/stream/dim_info.h | 14 +- src/stream/layouts.h | 7 +- src/types.stream.h | 4 + tests/CMakeLists.txt | 5 + tests/test_bench_input.c | 20 +-- tests/test_dimension.c | 7 +- tests/test_packed_input.c | 237 ++++++++++++++++++++++++++++++ tests/test_zarr_readback.c | 4 +- tests/test_zarr_readback_gpu.c | 4 +- 34 files changed, 570 insertions(+), 152 deletions(-) create mode 100644 tests/test_packed_input.c diff --git a/README.md b/README.md index 6ecf715e..3b141cb4 100644 --- a/README.md +++ b/README.md @@ -224,7 +224,8 @@ and multiscale, generated/image inputs, and discard/filesystem/S3/throttled sink measured input, and `--duration S` requests a minimum measured append time; neither changes that layout. The underlying frame dimension is unbounded. Source preparation and stream creation precede warmup. Generated inputs use a fixed 64 MiB source ring. Images use a preloaded, -chunk-padded ring; loading and padding are excluded from timing. +tightly packed ring; loading is excluded from timing. Chunk padding is handled +by the stream's scatter and LOD mapping. Warmup continues to a full batch boundary with no partial append-downsample accumulator, drains earlier work and sink metadata, and resets stage timings, @@ -240,8 +241,10 @@ TTY and top-level JSON rates/stage metrics describe this same window. The `measurement` object records the policy, requested and actual durations, warmup and measured work, effective per-dimension geometry, epoch/batch/staging sizes, and boundary samples. `throughput_in_gibs` counts submitted bytes; -`throughput_logical_gibs` excludes image padding. `compression_fold` counts -full decoded chunks, while `logical_compression_fold` counts logical input. +`throughput_logical_gibs` counts the same packed image bytes. `compression_fold` +counts full decoded chunks, while `logical_compression_fold` counts logical input. +Image results with `source_input_bytes` use packed input; older +`source_padded_bytes` results include caller padding in submitted-byte rates. Version 12 adds images to the common version-11 window. Neither is directly comparable to older whole-run or no-drain sustained rates. The specialized two-stream driver retains its explicit fixed-frame, whole-run policy and rejects timing options. @@ -394,7 +397,7 @@ error and unaccepted remainder before recycling the entire camera buffer. available batch. Both byte limits must respect the stream's input granularity, and the drain limit cannot exceed capacity. Zero drain limit uses capacity. -Size input capacity from the bytes actually submitted, including padding: +Size input capacity from the tightly packed bytes actually submitted: `capacity >= submitted_bytes_per_second * pause_seconds + active_drain_bytes` diff --git a/bench/bench_input.c b/bench/bench_input.c index df482f2a..2d02ae45 100644 --- a/bench/bench_input.c +++ b/bench/bench_input.c @@ -18,8 +18,6 @@ bench_input_load(struct bench_input* input, if (!input || !path || (dtype != dtype_u8 && dtype != dtype_u16 && dtype != dtype_f32) || !width || !height || !chunk_width || !chunk_height || - width > SIZE_MAX - (chunk_width - 1) || - height > SIZE_MAX - (chunk_height - 1) || width > SIZE_MAX / bpe / height || (bpe > 1 && *(const unsigned char*)&endian != 1)) { fprintf(stderr, @@ -27,13 +25,7 @@ bench_input_load(struct bench_input* input, return 1; } *input = (struct bench_input){ 0 }; - size_t padded_width = (width + chunk_width - 1) / chunk_width * chunk_width; - size_t padded_height = - (height + chunk_height - 1) / chunk_height * chunk_height; - if (padded_width > SIZE_MAX / bpe / padded_height) - return 1; size_t frame_elements = width * height; - size_t padded_frame = padded_width * padded_height; struct platform_clock clock = { 0 }; platform_toc(&clock); FILE* file = fopen(path, "rb"); @@ -49,28 +41,13 @@ bench_input_load(struct bench_input* input, fseek(file, 0, SEEK_SET)) goto Fail; size_t bytes = (size_t)length; - size_t frames = bytes / (frame_elements * bpe); - if (frames > SIZE_MAX / bpe / padded_frame) - goto Fail; - input->elements = frames * padded_frame; - input->frame_elements = padded_frame; + input->elements = bytes / bpe; + input->frame_elements = frame_elements; input->logical_frame_elements = frame_elements; input->source_bytes = bytes; - input->data = calloc(input->elements, bpe); - if (!input->data) + input->data = malloc(bytes); + if (!input->data || fread(input->data, 1, bytes, file) != bytes) goto Fail; - if (padded_frame == frame_elements) { - if (fread(input->data, 1, bytes, file) != bytes) - goto Fail; - } else { - for (size_t frame = 0; frame < frames; ++frame) - for (size_t row = 0; row < height; ++row) { - unsigned char* data = - input->data + (frame * padded_frame + row * padded_width) * bpe; - if (fread(data, bpe, width, file) != width) - goto Fail; - } - } if (fgetc(file) != EOF || ferror(file)) goto Fail; if (fclose(file)) { diff --git a/bench/bench_input.h b/bench/bench_input.h index 055c9495..cf07697a 100644 --- a/bench/bench_input.h +++ b/bench/bench_input.h @@ -16,8 +16,8 @@ struct bench_input float load_s; }; -// Prepared cyclic input. All offsets and sizes are writer bytes; logical -// frame bytes exclude image edge padding. Preparation precedes measurement. +// Prepared cyclic input. Frames are tightly packed; chunk padding is supplied +// by the stream mapping. All offsets and sizes are writer bytes. struct bench_source { const unsigned char* data; diff --git a/bench/bench_report.c b/bench/bench_report.c index eb4d3186..b885c7d2 100644 --- a/bench/bench_report.c +++ b/bench/bench_report.c @@ -1031,7 +1031,7 @@ print_bench_json_pass(const struct stream_metrics* m, jw_uint(&jw, measurement->append_bytes / dtype_bpe(dtype)); jw_key(&jw, "source_bytes"); jw_uint(&jw, images->input->source_bytes); - jw_key(&jw, "source_padded_bytes"); + jw_key(&jw, "source_input_bytes"); jw_uint(&jw, images->input->elements * dtype_bpe(dtype)); jw_key(&jw, "order"); jw_string(&jw, "cyclic"); diff --git a/bench/bench_util.c b/bench/bench_util.c index 89f29dae..4a7a1d94 100644 --- a/bench/bench_util.c +++ b/bench/bench_util.c @@ -1346,14 +1346,7 @@ bench_stream_main(int ac, char* av[], struct bench_spec spec) const size_t chunk_height = (size_t)dims[1].chunk_size; const size_t chunk_width = (size_t)dims[2].chunk_size; - size_t padded_width = - (size_t)((a.width + chunk_width - 1) / chunk_width * chunk_width); - size_t padded_height = - (size_t)((a.height + chunk_height - 1) / chunk_height * chunk_height); - if (padded_width > SIZE_MAX / bpe / padded_height) - return bench_failed(a.json_output); - size_t padded_frame = padded_width * padded_height; - if (a.frames > SIZE_MAX / bpe / padded_frame || + if (a.frames > SIZE_MAX / frame_bytes || a.append_elements > SIZE_MAX / bpe) { fprintf(stderr, "Image stream size overflows addressable memory\n"); return bench_failed(a.json_output); @@ -1364,7 +1357,7 @@ bench_stream_main(int ac, char* av[], struct bench_spec spec) if (!a.max_threads) a.max_threads = 4; if (!a.append_elements) - a.append_elements = padded_frame; + a.append_elements = frame_elements; if (bench_input_load(&input, a.input_path, a.dtype, diff --git a/bench/test_bench_images.py b/bench/test_bench_images.py index 5c640c9d..9a63b514 100644 --- a/bench/test_bench_images.py +++ b/bench/test_bench_images.py @@ -1,4 +1,4 @@ -"""Image replay shares coverage and reports logical/padded work independently.""" +"""Image replay submits packed pixels and pads only the output chunks.""" from array import array import json @@ -46,20 +46,18 @@ def check_type(dtype, typecode, replay_dtype, zarr_dtype): geometry = window["geometry"] assert replay["dtype"] == replay_dtype assert math.prod(replay["chunk_shape"]) * source.itemsize == 32 * 1024 - cy, cx = replay["chunk_shape"][1:] - padded_frame = math.ceil(height / cy) * cy * math.ceil(width / cx) * cx measured = replay["shape"][0] assert measured >= 7 - assert result["submitted_bytes"] == window["input_bytes"] == measured * padded_frame * source.itemsize + assert result["submitted_bytes"] == window["input_bytes"] == measured * frame * source.itemsize assert result["logical_input_bytes"] == measured * frame * source.itemsize - assert result["logical_input_bytes"] < result["submitted_bytes"] + assert result["logical_input_bytes"] == result["submitted_bytes"] assert math.isclose(result["throughput_in_gibs"], result["submitted_bytes"] / 2**30 / window["elapsed_s"], rel_tol=1e-5) assert math.isclose(result["throughput_logical_gibs"], result["logical_input_bytes"] / 2**30 / window["elapsed_s"], rel_tol=1e-5) assert replay["source_bytes"] == len(source) * source.itemsize - assert window["source_bytes"] == planes * padded_frame * source.itemsize - assert replay["source_padded_bytes"] == window["source_bytes"] + assert window["source_bytes"] == planes * frame * source.itemsize + assert replay["source_input_bytes"] == window["source_bytes"] assert replay["reference_shape"] == [64, height, width] # Compression keeps this filesystem metadata/index readback small. @@ -71,7 +69,7 @@ def check_type(dtype, typecode, replay_dtype, zarr_dtype): result = json.loads(process.stdout) window = result["measurement"] metadata = json.loads((output / "images" / "zarr.json").read_text()) - total_frames = (window["warmup_input_bytes"] + window["input_bytes"]) // (padded_frame * source.itemsize) + total_frames = (window["warmup_input_bytes"] + window["input_bytes"]) // (frame * source.itemsize) assert metadata["shape"] == [total_frames, height, width] assert metadata["data_type"] == zarr_dtype assert window["boundary_timing"] is False diff --git a/docs/buffered-writer.md b/docs/buffered-writer.md index 4d8b6814..5ed0b73d 100644 --- a/docs/buffered-writer.md +++ b/docs/buffered-writer.md @@ -11,15 +11,14 @@ sustained throughput. downstream batch; zero uses the full capacity. Capacity must be nonzero and the cap must not exceed it. Both values and input sizes must respect downstream granularity, such as whole elements. Capacity should cover the submitted byte rate -(including padding) times the pause allowance, plus input already occupied by an -active drain. Round up for whole arriving frames and downstream granularity. Keep +(excluding output chunk padding) times the pause allowance, plus input occupied +by an active drain. Round up for whole arriving frames and downstream granularity. Keep the drain limit independent of capacity so a larger queue does not create larger handoffs. The queue must have room when the pause begins; sustained overload or repeated pauses without recovery can still fill it. -For example, 4 useful GiB/s of BBBC022 frames becomes 4.23 GiB/s after padding -520×696 frames to 544×704 for 16 KiB chunks. Allowing 250 ms plus an active -64 MiB drain requires 1,152 MiB when rounded up to +For example, 4 GiB/s of tightly packed frames with a 250 ms pause allowance and an +active 64 MiB drain requires 1,088 MiB when rounded up to 64 MiB blocks. The camera pool is additional. Tune the sizes below to your workload: ```c diff --git a/scripts/datasets/README.md b/scripts/datasets/README.md index 26bd543b..57576604 100644 --- a/scripts/datasets/README.md +++ b/scripts/datasets/README.md @@ -186,9 +186,10 @@ Image chunks preserve the default T:Y:X bit ratio `1:4:4`; every requested decoded size is checked against the actual layout. All tiers retain the 0.5 GiB uncompressed full-shard floor, 1 GiB ceiling, 64 MiB batch target, four workers, one scale, and discard sink. Images retain native dimensions. The replay buffer -is zero-padded to whole chunks before timing; source and padded bytes are -reported separately. Primary throughput counts submitted bytes including this -padding; `throughput_logical_gibs` counts only native image bytes. A memory constraint or unattainable chunk target +contains tightly packed pixels. Scatter and LOD mapping pad output chunks +internally, so both input throughput fields count native image bytes. +`source_input_bytes` records the replay buffer size; archived runs used +`source_padded_bytes` for a caller-padded buffer. A memory constraint or unattainable chunk target causes failure instead of silently changing geometry. The unified runner writes the normal diff --git a/scripts/datasets/run.py b/scripts/datasets/run.py index f44ab6a4..ab4c8730 100644 --- a/scripts/datasets/run.py +++ b/scripts/datasets/run.py @@ -280,8 +280,12 @@ def check_result( if result["padded_input_bytes"] != padded_frames * padded_frame_bytes: raise ValueError("Padded input byte count is wrong") source_frames = pack["bytes"] // (pack["height"] * pack["width"] * 2) - if replay["source_padded_bytes"] != source_frames * padded_frame_bytes: - raise ValueError("Preloaded padded source byte count is wrong") + if "source_input_bytes" in replay: + source_valid = replay["source_input_bytes"] == pack["bytes"] + else: + source_valid = replay["source_padded_bytes"] == source_frames * padded_frame_bytes + if not source_valid: + raise ValueError("Preloaded source byte count is wrong") full_shard_bytes = ( math.prod(replay["chunk_shape"]) * math.prod(replay["chunks_per_shard"]) diff --git a/scripts/datasets/verify_output.py b/scripts/datasets/verify_output.py index 94a844c0..a9fe1d2c 100644 --- a/scripts/datasets/verify_output.py +++ b/scripts/datasets/verify_output.py @@ -57,25 +57,22 @@ def check_replay(report, images, minimum_frames, backend, profile): raise ValueError("Invalid image chunk shape") if math.prod(chunk) * images.dtype.itemsize != 32 * 1024: raise ValueError("Image chunk target changed") - padded_frame = math.prod( - (size + step - 1) // step * step - for size, step in zip((height, width), chunk[1:]) - ) * images.dtype.itemsize - measured, partial = divmod(window["input_bytes"], padded_frame) - warmup, warmup_partial = divmod(window["warmup_input_bytes"], padded_frame) + frame_bytes = height * width * images.dtype.itemsize + measured, partial = divmod(window["input_bytes"], frame_bytes) + warmup, warmup_partial = divmod(window["warmup_input_bytes"], frame_bytes) if partial or warmup_partial or measured < minimum_frames: raise ValueError("Image replay has partial or insufficient frames") expected = { "backend": backend, "codec": profile, "dtype": image_dtype(images), "shape": [measured, height, width], "source_bytes": images.nbytes, - "source_padded_bytes": len(images) * padded_frame, "order": "cyclic", + "source_input_bytes": len(images) * frame_bytes, "order": "cyclic", "codec_level": 0 if profile == "none" else 3, "shuffle": "bit" if profile.startswith("blosc-") else "none", } for key, value in expected.items(): if replay.get(key) != value: raise ValueError(f"Image replay changed {key}: expected {value}") - if (report["submitted_bytes"] != measured * padded_frame + if (report["submitted_bytes"] != measured * frame_bytes or report["logical_input_bytes"] != measured * height * width * images.dtype.itemsize): raise ValueError("Image logical/submitted byte accounting disagrees") return warmup + measured diff --git a/scripts/sweep/microscopy_data.py b/scripts/sweep/microscopy_data.py index 3d80a5b8..1d05a7b5 100644 --- a/scripts/sweep/microscopy_data.py +++ b/scripts/sweep/microscopy_data.py @@ -79,11 +79,15 @@ def check_observation(record, task, config, definition, resources=None): raise ValueError("Invalid logical byte accounting") frames, height, width = replay["shape"] frame_bytes = math.ceil(height / chunk[1]) * chunk[1] * math.ceil(width / chunk[2]) * chunk[2] * element_size + # Archived observations used caller-padded frames; new runs submit pixels. + if "source_input_bytes" in replay: + frame_bytes = height * width * element_size + source_bytes = replay.get("source_input_bytes", replay.get("source_padded_bytes")) if (window["input_bytes"] != frames * frame_bytes or result["submitted_bytes"] != window["input_bytes"] - or replay["source_padded_bytes"] != window["source_bytes"] + or source_bytes != window["source_bytes"] or replay["append_elements"] * element_size != window["append_bytes"]): - raise ValueError("Invalid padded byte accounting or replay buffer") + raise ValueError("Invalid input byte accounting or replay buffer") if definition["geometry_frames"] is not None and window["reference_frames"] != definition["geometry_frames"]: raise ValueError("Observation changed reference geometry") if not math.isfinite(result["process_wall_s"]) or result["process_wall_s"] <= 0: diff --git a/scripts/sweep/sweep.py b/scripts/sweep/sweep.py index 8e3be295..cffb7f29 100644 --- a/scripts/sweep/sweep.py +++ b/scripts/sweep/sweep.py @@ -588,17 +588,17 @@ def check_image_result( or (spec.chunk_depth is not None and chunk[0] != spec.chunk_depth)): raise ValueError("Image chunk geometry disagrees with requested target") height, width = pack["height"], pack["width"] - padded_frame = math.ceil(height / chunk[1]) * chunk[1] * math.ceil(width / chunk[2]) * chunk[2] * bytes_per_element - measured_frames, partial = divmod(window["input_bytes"], padded_frame) + frame_bytes = height * width * bytes_per_element + measured_frames, partial = divmod(window["input_bytes"], frame_bytes) if partial or measured_frames < frames: raise ValueError("Image measurement has partial or insufficient measured frames") if (replay["shape"] != [measured_frames, height, width] - or result["submitted_bytes"] != measured_frames * padded_frame + or result["submitted_bytes"] != measured_frames * frame_bytes or result["logical_input_bytes"] != measured_frames * height * width * bytes_per_element): raise ValueError("Image logical/submitted byte accounting disagrees") - source_bytes = len(pack["planes"]) * padded_frame - if replay["source_padded_bytes"] != source_bytes or window["source_bytes"] != source_bytes: - raise ValueError("Image padded source byte accounting disagrees") + source_bytes = len(pack["planes"]) * frame_bytes + if replay["source_input_bytes"] != source_bytes or window["source_bytes"] != source_bytes: + raise ValueError("Image packed source byte accounting disagrees") if type(result.get("worker_threads")) is not int or result["worker_threads"] != worker_threads: raise ValueError(f"Image benchmark did not use {worker_threads} workers") if spec.codec.startswith("blosc-") and result["blosc_block_bytes"] != spec.blosc_block_bytes: diff --git a/scripts/sweep/test_microscopy.py b/scripts/sweep/test_microscopy.py index 5f4304c3..7343e009 100644 --- a/scripts/sweep/test_microscopy.py +++ b/scripts/sweep/test_microscopy.py @@ -139,6 +139,24 @@ def test_unknown_options_and_schedule_tampering_fail(self): class StudyDataTests(unittest.TestCase): + def test_packed_input_accounting_and_archived_padding(self): + document = fixture() + # The original fixture retains the previous caller-padded contract. + validate_study(document) + for record in document["records"]: + result = record["result"] + replay, window = result["image_replay"], result["measurement"] + replay["source_input_bytes"] = replay["source_bytes"] + del replay["source_padded_bytes"] + window["source_bytes"] = replay["source_input_bytes"] + window["input_bytes"] = result["logical_input_bytes"] + result["submitted_bytes"] = result["logical_input_bytes"] + result["throughput_in_gibs"] = result["throughput_logical_gibs"] + validate_study(document) + document["records"][0]["result"]["measurement"]["source_bytes"] += 2 + with self.assertRaises(ValueError): + validate_study(document) + def test_logical_median_fold_and_observed_detail(self): document = fixture() case_id = next(task["case_id"] for task in document["plan"]["schedule"] diff --git a/src/cpu/stream.body.c b/src/cpu/stream.body.c index 8849844d..43a231ac 100644 --- a/src/cpu/stream.body.c +++ b/src/cpu/stream.body.c @@ -250,14 +250,33 @@ cpu_stream_flush_body(struct cpu_stream_view* v) // Flush partial epoch into the batch. if (*v->cursor_elements % v->layout->epoch_elements != 0) { + CHECK(Fail, *v->batch_accumulated < v->cl->epochs_per_batch); uint32_t active_mask = 1; + const size_t bpe = dtype_bpe(v->config->dtype); + const uint64_t used = *v->cursor_elements % v->layout->epoch_elements; if (v->levels->enable_multiscale) { + // A reused linear epoch may still contain the preceding epoch's tail. + memset((char*)v->linear + used * bpe, + 0, + (v->layout->epoch_elements - used) * bpe); struct scatter_epoch_params sp = make_scatter_params(v); CHECK(Fail, cpu_pipeline_scatter_epoch( &sp, *v->batch_accumulated, &active_mask) == 0); + } else { + void* epoch_pool = + (char*)v->chunk_pool + (uint64_t)*v->batch_accumulated * + v->levels->total_chunks * + v->layout->chunk_stride * bpe; + CHECK(Fail, + transpose_cpu(epoch_pool, + NULL, + (v->layout->epoch_elements - used) * bpe, + (uint8_t)bpe, + used, + v->layout, + v->pool) == 0); } - CHECK(Fail, *v->batch_accumulated < v->cl->epochs_per_batch); v->batch_active_masks[*v->batch_accumulated] = active_mask; (*v->batch_accumulated)++; } diff --git a/src/cpu/stream.c b/src/cpu/stream.c index 7f099a57..15e0f028 100644 --- a/src/cpu/stream.c +++ b/src/cpu/stream.c @@ -289,11 +289,10 @@ tile_stream_cpu_create(const struct tile_stream_configuration* config, // detect the at-capacity case without recomputing each call. { const struct dimension* dims = config->dimensions; - const uint8_t na = dim_info_n_append(&s->cl.dims); if (dims[0].size > 0) { - s->total_element_limit = s->layout.epoch_elements; - for (int d = 0; d < na; ++d) - s->total_element_limit *= ceildiv(dims[d].size, dims[d].chunk_size); + s->total_element_limit = 1; + for (int d = 0; d < config->rank; ++d) + s->total_element_limit *= dims[d].size; } else { s->total_element_limit = 0; } diff --git a/src/cpu/transpose.cpp b/src/cpu/transpose.cpp index dd96cc6e..c1e19a4c 100644 --- a/src/cpu/transpose.cpp +++ b/src/cpu/transpose.cpp @@ -47,8 +47,57 @@ struct transpose_ctx int64_t inner_stride; uint64_t epoch_elements; uint8_t bpe; + const struct tile_stream_layout* layout; }; +// Decode logical coordinates once per run ending at a row or chunk edge. +// The destination pool is already zeroed; cells outside input_shape are never +// visited. No padded frame or per-element division is needed. +static void +transpose_packed_range(size_t beg, size_t end, int tid, void* vctx) +{ + (void)tid; + const auto* c = (const transpose_ctx*)vctx; + const auto* l = c->layout; + const int rank = l->lifted_rank / 2; + while (beg < end) { + uint64_t rest = (c->i_offset + beg) % l->epoch_elements; + uint64_t count = end - beg; + int64_t out = 0; + for (int d = rank - 1; d >= 0; --d) { + const uint64_t coord = rest % l->input_shape[d]; + rest /= l->input_shape[d]; + const uint64_t chunk = l->lifted_shape[2 * d + 1]; + out += (coord / chunk) * l->lifted_strides[2 * d] + + (coord % chunk) * l->lifted_strides[2 * d + 1]; + if (d == rank - 1) { + const uint64_t row_left = l->input_shape[d] - coord; + const uint64_t chunk_left = chunk - coord % chunk; + if (count > row_left) + count = row_left; + if (count > chunk_left) + count = chunk_left; + } + } + char* dst = (char*)c->dst + out * c->bpe; + const char* src = c->src ? c->src + beg * c->bpe : nullptr; + const int64_t step = l->lifted_strides[l->lifted_rank - 1] * c->bpe; + if (step == c->bpe) { + if (src) + memcpy(dst, src, count * c->bpe); + else + memset(dst, 0, count * c->bpe); + } else + for (uint64_t i = 0; i < count; ++i) { + if (src) + memcpy(dst + i * step, src + i * c->bpe, c->bpe); + else + memset(dst + i * step, 0, c->bpe); + } + beg += count; + } +} + // The append chunk coordinates have zero strides: the caller selects their // destination epoch. All remaining coordinates describe one epoch. If their // strides are row-major, every range within the epoch can be copied directly. @@ -145,14 +194,17 @@ transpose_cpu(void* dst, inner_stride, layout->epoch_elements, bpe, + layout, }; // For small appends, dispatching and joining the pool costs more than the // scatter. Keep this choice independent of compression parallelism. constexpr uint64_t min_parallel_bytes = 64u << 10; + auto range = layout->has_partial_chunks || !src ? transpose_packed_range + : transpose_range; if (src_bytes < min_parallel_bytes) - transpose_range(0, n, 0, &c); + range(0, n, 0, &c); else - threadpool_for_n(pool, n, transpose_range, &c); + threadpool_for_n(pool, n, range, &c); return 0; } diff --git a/src/cpu/transpose.h b/src/cpu/transpose.h index 72305c17..f9cc34f5 100644 --- a/src/cpu/transpose.h +++ b/src/cpu/transpose.h @@ -13,8 +13,9 @@ extern "C" // CPU scatter transpose using the vadd() algorithm. // Copies directly when the layout proves that the epoch is contiguous; - // otherwise scatters src_bytes/bpe elements using its lifted shape/strides. + // otherwise maps packed logical coordinates into padded output chunks. // i_offset is the global flat input offset (for multi-call accumulation). + // A NULL src zeroes the mapped range (used for a final partial epoch). // Returns 0 on success and nonzero for an unsupported element size. int transpose_cpu(void* dst, const void* src, diff --git a/src/gpu/schedule.c b/src/gpu/schedule.c index 349b3dfb..39038806 100644 --- a/src/gpu/schedule.c +++ b/src/gpu/schedule.c @@ -773,10 +773,21 @@ schedule_add_partial_epoch(struct stream_engine* e, struct stream_context* ctx) { if (e->sched.accumulated >= e->sched.epochs_per_batch) return 1; + if (ctx->levels.enable_multiscale) { + const size_t bpe = dtype_bpe(ctx->config.dtype); + const uint64_t used = ctx->cursor_elements % ctx->layout.epoch_elements; + CU(Error, + cuMemsetD8Async(e->lod_shared.d_linear + used * bpe, + 0, + (ctx->layout.epoch_elements - used) * bpe, + e->streams.compute)); + } if (run_epoch_lod(e, ctx)) return 1; e->sched.accumulated++; return 0; +Error: + return 1; } // A worker left running owns shard state the caller reads and buffers destroy diff --git a/src/gpu/stream.ingest.c b/src/gpu/stream.ingest.c index 1420205b..4e418216 100644 --- a/src/gpu/stream.ingest.c +++ b/src/gpu/stream.ingest.c @@ -170,17 +170,14 @@ scatter_with_timing(struct staging_state* stage, struct scatter_timing* st = timing_begin(stage, bytes, compute); CHECK_SILENT(Fail, st); CHECK_SILENT(Fail, - transpose(gpu_pool_view_d(dst.first_epoch), - gpu_pool_view_d(d_in), - bytes, - (uint8_t)bpe, - in_epoch, - layout->epoch_elements, - dst.epoch_bytes, - layout->lifted_rank, - layout->lifted_shape, - layout->lifted_strides, - compute) == 0); + transpose_chunks(gpu_pool_view_d(dst.first_epoch), + gpu_pool_view_d(d_in), + bytes, + (uint8_t)bpe, + in_epoch, + dst.epoch_bytes, + layout, + compute) == 0); return timing_end(st, compute); Fail: diff --git a/src/gpu/stream.init.c b/src/gpu/stream.init.c index 86a557d3..4492c23c 100644 --- a/src/gpu/stream.init.c +++ b/src/gpu/stream.init.c @@ -243,14 +243,14 @@ engine_array_state_init(struct engine_array_state* st, for (int lv = 0; lv < cl->levels.nlod; ++lv) st->lod.layouts[lv] = cl->layouts[lv]; - // Every level scatters, so a layout the scatter cannot place is refused - // here rather than on the first append that reaches it. + // The lifted fast path needs trailing extents spanning exactly one epoch. + // Partial chunks use explicit logical coordinates in transpose_chunks. for (int lv = 0; lv < cl->levels.nlod; ++lv) { const struct tile_stream_layout* l = &st->lod.layouts[lv]; - if (transpose_check_layout(l->epoch_elements, - l->lifted_rank, - l->lifted_shape, - l->lifted_strides)) { + if (!l->has_partial_chunks && transpose_check_layout(l->epoch_elements, + l->lifted_rank, + l->lifted_shape, + l->lifted_strides)) { log_error("level %d has a layout the scatter cannot place", lv); goto Fail; } @@ -275,11 +275,10 @@ engine_array_state_init(struct engine_array_state* st, // append body detect the at-capacity case without recomputing each call. { const struct dimension* dims = ctx->config.dimensions; - const uint8_t na = dim_info_n_append(&ctx->dims); if (dims[0].size > 0) { - ctx->total_element_limit = ctx->layout.epoch_elements; - for (int d = 0; d < na; ++d) - ctx->total_element_limit *= ceildiv(dims[d].size, dims[d].chunk_size); + ctx->total_element_limit = 1; + for (int d = 0; d < ctx->config.rank; ++d) + ctx->total_element_limit *= dims[d].size; } } diff --git a/src/gpu/transpose.cu b/src/gpu/transpose.cu index bbb0a9ac..1ccba1a0 100644 --- a/src/gpu/transpose.cu +++ b/src/gpu/transpose.cu @@ -1,6 +1,7 @@ #include "defs.limits.h" #include "gpu/prelude.cuda.h" #include "gpu/transpose.h" +#include "stream/layouts.h" #include "util/prelude.h" #include @@ -10,8 +11,19 @@ constexpr int ELEMENTS_PER_BLOCK = (1 << 12) / (int)sizeof(T); // Travels in the launch parameters, so the kernel reads it from the constant // bank instead of paying a load per dimension per element. +template +struct scatter_edges +{}; + template -struct scatter_layout +struct scatter_edges +{ + Index chunks[HALF_MAX_RANK]; + Index chunk_strides[HALF_MAX_RANK]; +}; + +template +struct scatter_layout : scatter_edges { Index shape[MAX_RANK]; Index strides[MAX_RANK]; @@ -19,14 +31,14 @@ struct scatter_layout }; // Transpose data kernel - v0 -template +template __global__ void __launch_bounds__(256, 4) transpose_v0_k(T* d_dst, const T* d_src, Index src_size, Index i_offset, Index region_stride, - scatter_layout layout) + scatter_layout layout) { const Index block_offset = (Index)blockIdx.x * ELEMENTS_PER_BLOCK; const Index left = src_size - block_offset; @@ -40,7 +52,11 @@ __global__ void __launch_bounds__(256, 4) for (int d = layout.ndims - 1; d >= 0; --d) { const Index coord = rest % layout.shape[d]; rest /= layout.shape[d]; - out_offset += coord * layout.strides[d]; + if constexpr (Packed) + out_offset += (coord / layout.chunks[d]) * layout.chunk_strides[d] + + (coord % layout.chunks[d]) * layout.strides[d]; + else + out_offset += coord * layout.strides[d]; } // The visited extents span one epoch, so what is left counts epochs. @@ -60,6 +76,7 @@ struct transpose_args const uint64_t* shape; const int64_t* strides; CUstream stream; + const struct tile_stream_layout* packed = nullptr; }; template @@ -158,6 +175,17 @@ indices_fit_32_bits(const struct transpose_args& a, if (!add_within_32_bits(&max_offset, last / a.epoch_elements, region_stride)) return 0; + if (a.packed) { + for (int d = 0; d < a.rank; ++d) + if (a.shape[d] > UINT32_MAX || + a.packed->lifted_shape[2 * d + 1] > UINT32_MAX || + (uint64_t)a.strides[2 * d] > UINT32_MAX || + (uint64_t)a.strides[2 * d + 1] > UINT32_MAX) + return 0; + return add_within_32_bits( + &max_offset, 1, a.packed->chunks_per_epoch * a.packed->chunk_stride - 1); + } + for (int d = first_dim; d < a.rank; ++d) { if (a.shape[d] > UINT32_MAX || a.strides[d] < 0) return 0; @@ -168,25 +196,30 @@ indices_fit_32_bits(const struct transpose_args& a, return 1; } -template +template static int launch(const struct transpose_args& a, int first_dim) { const uint64_t src_size = a.src_bytes / sizeof(T); const uint64_t region_stride = a.region_bytes / sizeof(T); - scatter_layout layout = {}; + scatter_layout layout = {}; layout.ndims = a.rank - first_dim; for (int d = 0; d < layout.ndims; ++d) { layout.shape[d] = (Index)a.shape[first_dim + d]; - layout.strides[d] = (Index)a.strides[first_dim + d]; + if constexpr (Packed) { + layout.strides[d] = (Index)a.strides[2 * d + 1]; + layout.chunks[d] = (Index)a.packed->lifted_shape[2 * d + 1]; + layout.chunk_strides[d] = (Index)a.strides[2 * d]; + } else + layout.strides[d] = (Index)a.strides[first_dim + d]; } const int block_size = 256; const unsigned grid_size = (unsigned)ceildiv(src_size, (uint64_t)ELEMENTS_PER_BLOCK); - return CUDA_LAUNCH(transpose_v0_k + return CUDA_LAUNCH(transpose_v0_k <<>>( (T*)a.d_dst_beg, (const T*)a.d_src_beg, @@ -206,6 +239,12 @@ transpose_launch(const struct transpose_args& a) if (!args_valid(a)) return 1; + if (a.packed) { + if (indices_fit_32_bits(a, src_size, 0, a.region_bytes / sizeof(T))) + return launch(a, 0); + return launch(a, 0); + } + const int first_dim = first_decomposed_dim(a.epoch_elements, a.rank, a.shape, a.strides); if (first_dim < 0) @@ -216,6 +255,54 @@ transpose_launch(const struct transpose_args& a) return launch(a, first_dim); } +extern "C" int +transpose_chunks(CUdeviceptr dst, + CUdeviceptr src, + uint64_t src_bytes, + uint8_t bpe, + uint64_t i_offset, + uint64_t region_bytes, + const struct tile_stream_layout* layout, + CUstream stream) +{ + if (!layout->has_partial_chunks) + return transpose(dst, + src, + src_bytes, + bpe, + i_offset, + layout->epoch_elements, + region_bytes, + layout->lifted_rank, + layout->lifted_shape, + layout->lifted_strides, + stream); + + const struct transpose_args a = { dst, + src, + src_bytes, + i_offset, + layout->epoch_elements, + region_bytes, + (uint8_t)(layout->lifted_rank / 2), + layout->input_shape, + layout->lifted_strides, + stream, + layout }; + switch (bpe) { + case 1: + return transpose_launch(a); + case 2: + return transpose_launch(a); + case 4: + return transpose_launch(a); + case 8: + return transpose_launch(a); + default: + return 1; + } +} + extern "C" int transpose(CUdeviceptr d_dst_beg, CUdeviceptr d_src_beg, diff --git a/src/gpu/transpose.h b/src/gpu/transpose.h index 0b7f8b9c..3d40b05d 100644 --- a/src/gpu/transpose.h +++ b/src/gpu/transpose.h @@ -10,6 +10,19 @@ extern "C" { #endif + struct tile_stream_layout; + + // Map packed acquisition coordinates directly into zeroed chunk pools. + // The range can cross epochs; region_bytes separates destination epochs. + int transpose_chunks(CUdeviceptr dst, + CUdeviceptr src, + uint64_t src_bytes, + uint8_t bpe, + uint64_t i_offset, + uint64_t region_bytes, + const struct tile_stream_layout* layout, + CUstream stream); + // The layout half of what transpose checks, so a stream can refuse a layout // it could never scatter when it opens rather than partway through an // append. Returns 0 when the scatter can place the layout. diff --git a/src/lod/lod_plan.c b/src/lod/lod_plan.c index 49d803d2..e8f20dfa 100644 --- a/src/lod/lod_plan.c +++ b/src/lod/lod_plan.c @@ -605,7 +605,9 @@ dims_print(const struct dimension* dims, uint8_t rank) if (i >= na) chunks_per_epoch *= tc; } - double epoch_elements = (double)chunks_per_epoch * (double)chunk_elements; + double epoch_elements = 1; + for (uint8_t i = 0; i < rank; ++i) + epoch_elements *= i < na ? dims[i].chunk_size : dims[i].size; fprintf(stderr, "chunk_elements: %llu chunks/epoch: %llu epoch_elements: %.3g\n", (unsigned long long)chunk_elements, diff --git a/src/multiarray/stream.c b/src/multiarray/stream.c index eff7294c..a385aaf0 100644 --- a/src/multiarray/stream.c +++ b/src/multiarray/stream.c @@ -195,11 +195,10 @@ init_array_descriptor(struct array_descriptor* desc, // total_element_limit: configured stream length (0 = unbounded) { const struct dimension* dims = config->dimensions; - const uint8_t na = dim_info_n_append(&desc->cl.dims); if (dims[0].size > 0) { - desc->total_element_limit = desc->layout.epoch_elements; - for (int d = 0; d < na; ++d) - desc->total_element_limit *= ceildiv(dims[d].size, dims[d].chunk_size); + desc->total_element_limit = 1; + for (int d = 0; d < config->rank; ++d) + desc->total_element_limit *= dims[d].size; } else { desc->total_element_limit = 0; } diff --git a/src/stream/config.c b/src/stream/config.c index e8dc3b7b..d26c9f71 100644 --- a/src/stream/config.c +++ b/src/stream/config.c @@ -89,6 +89,8 @@ compute_level_layout(struct tile_stream_layout* layout, layout->lifted_rank = 2 * rank; layout->chunk_elements = 1; + layout->epoch_elements = 1; + layout->has_partial_chunks = 0; uint64_t chunk_count[HALF_MAX_RANK]; for (int i = 0; i < rank; ++i) { @@ -97,6 +99,12 @@ compute_level_layout(struct tile_stream_layout* layout, layout->lifted_shape[2 * i] = chunk_count[i]; layout->lifted_shape[2 * i + 1] = dims[i].chunk_size; layout->chunk_elements *= dims[i].chunk_size; + const uint64_t extent = i < n_append ? dims[i].chunk_size : level_shape[i]; + CHECK(Fail, extent > 0); + layout->input_shape[i] = extent; + layout->has_partial_chunks |= extent % dims[i].chunk_size != 0; + CHECK_MUL_OVERFLOW(Fail, layout->epoch_elements, extent, UINT64_MAX); + layout->epoch_elements *= extent; } { @@ -123,7 +131,6 @@ compute_level_layout(struct tile_stream_layout* layout, layout->chunks_per_epoch *= chunk_count[d]; CHECK_MUL_OVERFLOW( Fail, layout->chunks_per_epoch, layout->chunk_elements, UINT64_MAX); - layout->epoch_elements = layout->chunks_per_epoch * layout->chunk_elements; // Collapse all append dims for (int d = 0; d < n_append; ++d) layout->lifted_strides[2 * d] = 0; diff --git a/src/stream/dim_info.h b/src/stream/dim_info.h index d90144c2..0227b9f7 100644 --- a/src/stream/dim_info.h +++ b/src/stream/dim_info.h @@ -83,12 +83,11 @@ dim_info_readable_append_sizes(const struct dim_info* info, // Compute exact append dim sizes for final metadata. // // Bounded append dims (1..n_append-1) report their declared size. -// Dim 0: derived from cursor_elements / product(the encoded inner extents). -// Inner extents are rounded to whole chunks because the stream cursor counts -// edge padding as input elements. +// Dim 0: derived from cursor_elements / product(the logical inner extents). +// The cursor counts tightly packed input, excluding output chunk padding. // For LOD level > 0 with append_downsample: ceildiv(dim0, 2^level). // -// Precondition: cursor_elements is a multiple of the padded inner-dim product +// Precondition: cursor_elements is a multiple of the logical inner-dim product // (no partial frames). static inline void dim_info_final_append_sizes(const struct dim_info* info, @@ -104,11 +103,8 @@ dim_info_final_append_sizes(const struct dim_info* info, for (const struct dimension* d = info->append.beg + 1; d < info->append.end; ++d) inner *= d->size; - for (const struct dimension* d = info->inner.beg; d < info->inner.end; ++d) { - const uint64_t chunks = - d->size / d->chunk_size + (d->size % d->chunk_size != 0); - inner *= chunks * d->chunk_size; - } + for (const struct dimension* d = info->inner.beg; d < info->inner.end; ++d) + inner *= d->size; uint64_t dim0 = inner > 0 ? cursor_elements / inner : 0; if (level > 0 && info->append_downsample) diff --git a/src/stream/layouts.h b/src/stream/layouts.h index 1527b7e8..20af9ab9 100644 --- a/src/stream/layouts.h +++ b/src/stream/layouts.h @@ -15,10 +15,15 @@ struct tile_stream_layout uint64_t lifted_shape[MAX_RANK]; int64_t lifted_strides[MAX_RANK]; + // Packed acquisition extents within one epoch. Append dimensions span one + // chunk; inner dimensions retain their logical size, including partial edges. + uint64_t input_shape[HALF_MAX_RANK]; + int has_partial_chunks; + uint64_t chunk_elements; uint64_t chunk_stride; uint64_t chunks_per_epoch; - uint64_t epoch_elements; + uint64_t epoch_elements; // logical input elements, excluding chunk padding size_t chunk_pool_bytes; }; diff --git a/src/types.stream.h b/src/types.stream.h index 8e34e18c..e20a5365 100644 --- a/src/types.stream.h +++ b/src/types.stream.h @@ -133,6 +133,10 @@ struct stream_metrics uint64_t memcpy_bytes; }; +// Streams accept tightly packed elements in acquisition dimension order. +// Appends may split rows or frames at element boundaries. Scatter/LOD mapping +// places them into chunks and zero-fills output edges; callers supply no +// chunk padding. Cursors and configured capacity count only logical elements. struct tile_stream_configuration { size_t buffer_capacity_bytes; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 40793a1b..e073339d 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -129,6 +129,11 @@ add_gpu_test_exe(test_transpose test_transpose.c LINKS CUDA::cuda_driver CUDA::c # CPU transpose test (no CUDA) add_test_exe(test_transpose_cpu test_transpose_cpu.c transpose_cpu threadpool index_ops_util chucky_log) +add_test_exe(test_packed_input_cpu test_packed_input.c stream_cpu test_shard_sink test_shard_verify) +add_gpu_test_exe(test_packed_input_gpu test_packed_input.c LINKS stream test_shard_sink test_shard_verify CUDA::cuda_driver CUDA::cudart_static) +if(CHUCKY_ENABLE_GPU) + target_compile_definitions(test_packed_input_gpu PRIVATE TEST_GPU) +endif() # CPU compress test (no CUDA) add_test_exe(test_compress_cpu test_compress_cpu.c compress_cpu threadpool chucky_log) diff --git a/tests/test_bench_input.c b/tests/test_bench_input.c index 86e1a546..8d5b0b5b 100644 --- a/tests/test_bench_input.c +++ b/tests/test_bench_input.c @@ -55,7 +55,7 @@ replay(const struct bench_source* source, } static int -test_load_padding(enum dtype dtype) +test_load_packed(enum dtype dtype) { char directory[1024]; char path[1200]; @@ -78,20 +78,12 @@ test_load_padding(enum dtype dtype) goto Cleanup; if (bench_input_load(&input, path, dtype, 65, 66, 64, 64) || - input.elements != 2 * 128 * 128 || input.frame_elements != 128 * 128 || + input.elements != 2 * 65 * 66 || input.frame_elements != 65 * 66 || input.logical_frame_elements != 65 * 66 || input.source_bytes != source_bytes) goto Cleanup; - for (size_t frame = 0; frame < 2; ++frame) - for (size_t y = 0; y < 128; ++y) - for (size_t x = 0; x < 128; ++x) - for (size_t byte = 0; byte < bpe; ++byte) { - unsigned char value = - y < 66 && x < 65 ? source[((frame * 66 + y) * 65 + x) * bpe + byte] - : 0; - if (input.data[((frame * 128 + y) * 128 + x) * bpe + byte] != value) - goto Cleanup; - } + if (memcmp(input.data, source, source_bytes)) + goto Cleanup; if (!bench_input_load(&input, path, dtype, SIZE_MAX, 1, 64, 64) || !bench_input_load(&input, path, dtype, 1, 1, 0, 64) || !bench_input_load(&input, path, dtype_f64, 65, 66, 64, 64)) @@ -160,8 +152,8 @@ test_generated_prefix(void) int main(void) { - if (test_generated_prefix() || test_load_padding(dtype_u8) || - test_load_padding(dtype_u16) || test_load_padding(dtype_f32)) + if (test_generated_prefix() || test_load_packed(dtype_u8) || + test_load_packed(dtype_u16) || test_load_packed(dtype_f32)) return 1; uint16_t data[7]; for (size_t i = 0; i < 7; ++i) diff --git a/tests/test_dimension.c b/tests/test_dimension.c index 89449eab..ddf21b46 100644 --- a/tests/test_dimension.c +++ b/tests/test_dimension.c @@ -913,9 +913,8 @@ test_dim_info_final_append_sizes(void) dim_info_final_append_sizes(&info, cursor, 1, append_sizes); CHECK(Error, append_sizes[0] == 12); - // The cursor includes edge padding for non-divisible inner dimensions. - // 600x600 is encoded as 640x640 with 64x64 chunks, but still reports 13 - // logical frames. + // Edge padding belongs to the output chunks, not the input cursor. + // 13 tightly packed 600x600 frames still report 13 with 64x64 chunks. { struct dimension padded_dims[3]; uint64_t padded_sizes[] = { 0, 600, 600 }; @@ -926,7 +925,7 @@ test_dim_info_final_append_sizes(void) struct dim_info padded_info; CHECK(Error, dim_info_init(&padded_info, padded_dims, 3) == 0); dim_info_final_append_sizes( - &padded_info, 13ull * 640 * 640, 0, append_sizes); + &padded_info, 13ull * 600 * 600, 0, append_sizes); CHECK(Error, append_sizes[0] == 13); } diff --git a/tests/test_packed_input.c b/tests/test_packed_input.c new file mode 100644 index 00000000..b54366a1 --- /dev/null +++ b/tests/test_packed_input.c @@ -0,0 +1,237 @@ +// Packed input must land at logical coordinates, leaving chunk edges zero. +// Compile the same regression for both stream backends. +#include "test_shard_sink.h" +#include "test_shard_verify.h" +#include "util/prelude.h" + +#ifdef TEST_GPU +#include "stream.gpu.h" +#include "test_runner.h" +#define stream_create tile_stream_gpu_create +#define stream_writer tile_stream_gpu_writer +#define stream_cursor tile_stream_gpu_cursor +#define stream_destroy tile_stream_gpu_destroy +#define stream_type tile_stream_gpu +#else +#include "stream.cpu.h" +#define stream_create tile_stream_cpu_create +#define stream_writer tile_stream_cpu_writer +#define stream_cursor tile_stream_cpu_cursor +#define stream_destroy tile_stream_cpu_destroy +#define stream_type tile_stream_cpu +#endif + +#include +#include + +static uint16_t +pixel(uint64_t t, uint64_t y, uint64_t x, uint64_t width) +{ + return (uint16_t)(1 + 1000 * t + y * width + x); +} + +static int +run_case(uint64_t frames, + uint64_t height, + uint64_t width, + uint64_t depth, + size_t append_elements, + int bounded, + int transpose, + int multiscale, + uint64_t channels) +{ + int error = 1; + struct stream_type* stream = NULL; + struct test_shard_sink sink; + log_info("packed: %llux%llux%llux%llu depth=%llu append=%zu bounded=%d " + "transpose=%d lod=%d", + (unsigned long long)frames, + (unsigned long long)channels, + (unsigned long long)height, + (unsigned long long)width, + (unsigned long long)depth, + append_elements, + bounded, + transpose, + multiscale); + const int shard_counts[] = { 8, 8 }; + const int nlod = multiscale ? 2 : 1; + test_sink_init_multi(&sink, nlod, shard_counts, 4 << 20); + uint16_t* input = + malloc((frames + 1) * channels * height * width * sizeof(*input)); + CHECK(Cleanup, input); + for (uint64_t t = 0; t < (frames + 1) * channels; ++t) + for (uint64_t y = 0; y < height; ++y) + for (uint64_t x = 0; x < width; ++x) + input[(t * height + y) * width + x] = pixel(t, y, x, width); + + struct dimension dims[4]; + const int rank = channels == 1 ? 3 : 4; + if (rank == 3) { + dims_create( + dims, "tyx", (uint64_t[]){ bounded ? frames : 0, height, width }); + dims_set_chunk_sizes(dims, 3, (uint64_t[]){ depth, 8, 8 }); + dims_set_shard_counts(dims, 3, (uint64_t[]){ 0, 1, 1 }); + } else { + dims_create(dims, + "tcyx", + (uint64_t[]){ bounded ? frames : 0, channels, height, width }); + dims_set_chunk_sizes(dims, 4, (uint64_t[]){ depth, 1, 8, 8 }); + dims_set_shard_counts(dims, 4, (uint64_t[]){ 0, 1, 1, 1 }); + } + dims[0].chunks_per_shard = 4; + dims[rank - 2].downsample = dims[rank - 1].downsample = multiscale; + if (transpose) { + dims[rank - 2].storage_position = rank - 1; + dims[rank - 1].storage_position = rank - 2; + } + const struct tile_stream_configuration config = { + .buffer_capacity_bytes = 510, + .dtype = dtype_u16, + .rank = rank, + .dimensions = dims, + .codec = { .id = CODEC_NONE }, + .epochs_per_batch = 2, + .max_nlod = nlod, + .max_threads = 4, + .reduce_method = lod_reduce_max, + }; + stream = stream_create(&config, &sink.base); + CHECK(Cleanup, stream); + struct writer* writer = stream_writer(stream); + const size_t total = frames * channels * height * width; + for (size_t offset = 0; offset < total;) { + size_t count = append_elements; + if (count > total - offset) + count = total - offset; + const struct slice part = { input + offset, input + offset + count }; + const struct writer_result r = writer_append_wait(writer, part); + CHECK(Cleanup, !r.error && r.rest.beg == r.rest.end); + offset += count; + } + if (bounded) { + const struct slice extra = { input + total, + input + total + height * width }; + const struct writer_result r = writer_append_wait(writer, extra); + CHECK(Cleanup, r.error == writer_error_finished && r.rest.beg == extra.beg); + } + CHECK(Cleanup, stream_cursor(stream) == total); + CHECK(Cleanup, !writer_flush(writer).error); + CHECK(Cleanup, !writer_close(writer).error); + CHECK(Cleanup, sink.last_append_size0 == frames); + + for (int lv = 0; lv < nlod; ++lv) { + const uint64_t scale = 1u << lv; + const uint64_t ny = ceildiv(height, scale), nx = ceildiv(width, scale); + const uint64_t cy = ceildiv(ny, 8), cx = ceildiv(nx, 8); + const uint64_t ct = + bounded && ceildiv(frames, depth) < 4 ? ceildiv(frames, depth) : 4; + const size_t slots = ct * channels * cy * cx; + uint64_t* offsets = malloc(slots * sizeof(*offsets)); + uint64_t* sizes = malloc(slots * sizeof(*sizes)); + if (!offsets || !sizes) { + free(offsets); + free(sizes); + goto Cleanup; + } + int bad = 0; + for (uint64_t sh = 0; sh < ceildiv(frames, ct * depth) && !bad; ++sh) { + const struct test_shard_writer* w = &sink.writers[lv][sh]; + if (!w->finalized || shard_index_check_crc(w->buf, w->size, slots) || + shard_index_parse(w->buf, w->size, slots, offsets, sizes)) { + bad = 1; + break; + } + for (uint64_t slot = 0; slot < slots && !bad; ++slot) { + const uint64_t tc = sh * ct + slot / (channels * cy * cx); + const uint64_t channel = slot / (cy * cx) % channels; + if (tc * depth >= frames) { + bad = sizes[slot] != UINT64_MAX; + continue; + } + if (sizes[slot] != depth * 8 * 8 * sizeof(uint16_t) || + offsets[slot] > w->size || sizes[slot] > w->size - offsets[slot]) { + bad = 1; + break; + } + const uint64_t yc = transpose ? slot % cy : (slot / cx) % cy; + const uint64_t xc = transpose ? (slot / cy) % cx : slot % cx; + for (uint64_t t = 0; t < depth; ++t) + for (uint64_t y = 0; y < 8; ++y) + for (uint64_t x = 0; x < 8; ++x) { + uint64_t gt = tc * depth + t, gy = yc * 8 + y, gx = xc * 8 + x; + uint16_t expected = 0, actual; + if (gt < frames && gy < ny && gx < nx) { + // Independent max-reduction reference, including odd edges. + for (uint64_t iy = gy * scale; + iy < (gy + 1) * scale && iy < height; + ++iy) + for (uint64_t ix = gx * scale; + ix < (gx + 1) * scale && ix < width; + ++ix) { + uint16_t value = + pixel(gt * channels + channel, iy, ix, width); + if (value > expected) + expected = value; + } + } + const uint64_t pos = t * 64 + (transpose ? x * 8 + y : y * 8 + x); + memcpy(&actual, w->buf + offsets[slot] + pos * 2, 2); + if (actual != expected && !bad) { + log_error("level %d pixel (%llu,%llu,%llu): %u != %u", + lv, + (unsigned long long)gt, + (unsigned long long)gy, + (unsigned long long)gx, + actual, + expected); + bad = 1; + } + } + } + } + free(offsets); + free(sizes); + CHECK(Cleanup, !bad); + } + error = 0; +Cleanup: + stream_destroy(stream); + test_sink_free(&sink); + free(input); + return error; +} + +static int +test_packed(void) +{ + // The reported reproducer, a divisible control, sub-chunk dimensions, + // split rows/frames, reused batches/shards, bounded tails and storage order. + CHECK(Fail, !run_case(3, 9, 13, 1, 117, 0, 0, 0, 1)); + CHECK(Fail, !run_case(3, 8, 16, 1, 128, 0, 0, 0, 1)); + CHECK(Fail, !run_case(3, 3, 5, 2, 7, 0, 0, 0, 1)); + for (int bounded = 0; bounded < 2; ++bounded) + for (int transpose = 0; transpose < 2; ++transpose) + for (int lod = 0; lod < 2; ++lod) + CHECK(Fail, !run_case(11, 9, 13, 2, 131, bounded, transpose, lod, 1)); + CHECK(Fail, !run_case(3, 257, 263, 1, 70001, 0, 0, 0, 1)); + CHECK(Fail, !run_case(11, 8, 16, 2, 131, 1, 0, 0, 1)); + for (int lod = 0; lod < 2; ++lod) { + CHECK(Fail, !run_case(3, 9, 13, 1, 131, 0, 1, lod, 3)); + CHECK(Fail, !run_case(3, 9, 13, 2, 131, 1, 1, lod, 3)); + } + return 0; +Fail: + return 1; +} + +#ifdef TEST_GPU +RUN_GPU_TESTS({ "packed_input", test_packed }) +#else +int +main(void) +{ + return test_packed(); +} +#endif diff --git a/tests/test_zarr_readback.c b/tests/test_zarr_readback.c index 61b67c15..72fcf2de 100644 --- a/tests/test_zarr_readback.c +++ b/tests/test_zarr_readback.c @@ -13,8 +13,8 @@ #include #define NT 4 -#define NY 8 -#define NX 12 +#define NY 9 +#define NX 13 // --- write_zarr --- diff --git a/tests/test_zarr_readback_gpu.c b/tests/test_zarr_readback_gpu.c index 3a4ee0e1..1012c3e4 100644 --- a/tests/test_zarr_readback_gpu.c +++ b/tests/test_zarr_readback_gpu.c @@ -12,8 +12,8 @@ #include #define NT 4 -#define NY 256 -#define NX 256 +#define NY 255 +#define NX 257 static int write_zarr(const char* store_path, struct codec_config codec, int buffered) From ae68b192a875771b87fb2a290a34c60f73943d85 Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Sun, 4 Oct 2026 10:28:18 -0700 Subject: [PATCH 2/4] Simplify packed input layouts and benchmark accounting --- bench/bench_input.c | 9 ++------ bench/bench_input.h | 7 +----- bench/bench_report.c | 2 +- bench/bench_util.c | 20 ++++------------- src/cpu/transpose.cpp | 5 +++-- src/gpu/stream.init.c | 9 ++++---- src/gpu/transpose.cu | 47 ++++++++++++++++++---------------------- src/stream/config.c | 2 -- src/stream/layouts.h | 9 +++++++- tests/test_bench_input.c | 16 ++++++-------- 10 files changed, 52 insertions(+), 74 deletions(-) diff --git a/bench/bench_input.c b/bench/bench_input.c index 2d02ae45..c608779f 100644 --- a/bench/bench_input.c +++ b/bench/bench_input.c @@ -9,16 +9,13 @@ bench_input_load(struct bench_input* input, const char* path, enum dtype dtype, size_t width, - size_t height, - size_t chunk_width, - size_t chunk_height) + size_t height) { const uint16_t endian = 1; const size_t bpe = dtype_bpe(dtype); if (!input || !path || (dtype != dtype_u8 && dtype != dtype_u16 && dtype != dtype_f32) || - !width || !height || !chunk_width || !chunk_height || - width > SIZE_MAX / bpe / height || + !width || !height || width > SIZE_MAX / bpe / height || (bpe > 1 && *(const unsigned char*)&endian != 1)) { fprintf(stderr, "Image replay requires little-endian u8, u16, or f32 frames\n"); @@ -41,9 +38,7 @@ bench_input_load(struct bench_input* input, fseek(file, 0, SEEK_SET)) goto Fail; size_t bytes = (size_t)length; - input->elements = bytes / bpe; input->frame_elements = frame_elements; - input->logical_frame_elements = frame_elements; input->source_bytes = bytes; input->data = malloc(bytes); if (!input->data || fread(input->data, 1, bytes, file) != bytes) diff --git a/bench/bench_input.h b/bench/bench_input.h index cf07697a..e04bb948 100644 --- a/bench/bench_input.h +++ b/bench/bench_input.h @@ -9,9 +9,7 @@ struct bench_input { unsigned char* data; - size_t elements; size_t frame_elements; - size_t logical_frame_elements; size_t source_bytes; float load_s; }; @@ -23,7 +21,6 @@ struct bench_source const unsigned char* data; size_t bytes; size_t frame_bytes; - size_t logical_frame_bytes; }; struct slice @@ -36,9 +33,7 @@ bench_input_load(struct bench_input* input, const char* path, enum dtype dtype, size_t width, - size_t height, - size_t chunk_width, - size_t chunk_height); + size_t height); void bench_input_free(struct bench_input* input); diff --git a/bench/bench_report.c b/bench/bench_report.c index b885c7d2..b7559462 100644 --- a/bench/bench_report.c +++ b/bench/bench_report.c @@ -1032,7 +1032,7 @@ print_bench_json_pass(const struct stream_metrics* m, jw_key(&jw, "source_bytes"); jw_uint(&jw, images->input->source_bytes); jw_key(&jw, "source_input_bytes"); - jw_uint(&jw, images->input->elements * dtype_bpe(dtype)); + jw_uint(&jw, images->input->source_bytes); jw_key(&jw, "order"); jw_string(&jw, "cyclic"); jw_key(&jw, "load_s"); diff --git a/bench/bench_util.c b/bench/bench_util.c index 4a7a1d94..6d2e7a30 100644 --- a/bench/bench_util.c +++ b/bench/bench_util.c @@ -599,9 +599,8 @@ run_bench(const struct bench_config* cfg) if (cfg->input) { source = (struct bench_source){ .data = (const unsigned char*)cfg->input->data, - .bytes = cfg->input->elements * bpe, + .bytes = cfg->input->source_bytes, .frame_bytes = cfg->input->frame_elements * bpe, - .logical_frame_bytes = cfg->input->logical_frame_elements * bpe, }; measurement.source_bytes = source.bytes; if (!cfg->append_elements) @@ -617,7 +616,6 @@ run_bench(const struct bench_config* cfg) source.frame_bytes = bpe; for (uint8_t d = 1; d < rank; ++d) source.frame_bytes *= dims[d].size; - source.logical_frame_bytes = source.frame_bytes; } measurement.prep_s = (platform_monotonic_ns() - prep_start) * 1e-9 + (cfg->input ? cfg->input->load_s : 0); @@ -850,10 +848,7 @@ run_bench(const struct bench_config* cfg) const float wall_s = (float)measurement.elapsed_s; measurement.output_bytes = sink_bytes(&meter, &tss, &dss) - measurement.warmup_output_bytes; - measurement.logical_input_bytes = cfg->input ? measurement.input_bytes / - source.frame_bytes * - source.logical_frame_bytes - : measurement.input_bytes; + measurement.logical_input_bytes = measurement.input_bytes; if (platform_peak_resident_memory(&mem_used.host_peak_bytes) != 0) { log_warn(" host peak memory reading unavailable"); @@ -1344,8 +1339,6 @@ bench_stream_main(int ac, char* av[], struct bench_spec spec) return bench_failed(a.json_output); } - const size_t chunk_height = (size_t)dims[1].chunk_size; - const size_t chunk_width = (size_t)dims[2].chunk_size; if (a.frames > SIZE_MAX / frame_bytes || a.append_elements > SIZE_MAX / bpe) { fprintf(stderr, "Image stream size overflows addressable memory\n"); @@ -1358,13 +1351,8 @@ bench_stream_main(int ac, char* av[], struct bench_spec spec) a.max_threads = 4; if (!a.append_elements) a.append_elements = frame_elements; - if (bench_input_load(&input, - a.input_path, - a.dtype, - (size_t)a.width, - (size_t)a.height, - chunk_width, - chunk_height)) + if (bench_input_load( + &input, a.input_path, a.dtype, (size_t)a.width, (size_t)a.height)) return bench_failed(a.json_output); } else if (a.input_path || a.width || a.height || a.chunk_depth) { fprintf(stderr, "Image options require bench_stream_microscopy\n"); diff --git a/src/cpu/transpose.cpp b/src/cpu/transpose.cpp index c1e19a4c..39eba67b 100644 --- a/src/cpu/transpose.cpp +++ b/src/cpu/transpose.cpp @@ -199,8 +199,9 @@ transpose_cpu(void* dst, // For small appends, dispatching and joining the pool costs more than the // scatter. Keep this choice independent of compression parallelism. constexpr uint64_t min_parallel_bytes = 64u << 10; - auto range = layout->has_partial_chunks || !src ? transpose_packed_range - : transpose_range; + auto range = layout_has_partial_chunks(layout) || !src + ? transpose_packed_range + : transpose_range; if (src_bytes < min_parallel_bytes) range(0, n, 0, &c); else diff --git a/src/gpu/stream.init.c b/src/gpu/stream.init.c index 4492c23c..61920450 100644 --- a/src/gpu/stream.init.c +++ b/src/gpu/stream.init.c @@ -247,10 +247,11 @@ engine_array_state_init(struct engine_array_state* st, // Partial chunks use explicit logical coordinates in transpose_chunks. for (int lv = 0; lv < cl->levels.nlod; ++lv) { const struct tile_stream_layout* l = &st->lod.layouts[lv]; - if (!l->has_partial_chunks && transpose_check_layout(l->epoch_elements, - l->lifted_rank, - l->lifted_shape, - l->lifted_strides)) { + if (!layout_has_partial_chunks(l) && + transpose_check_layout(l->epoch_elements, + l->lifted_rank, + l->lifted_shape, + l->lifted_strides)) { log_error("level %d has a layout the scatter cannot place", lv); goto Fail; } diff --git a/src/gpu/transpose.cu b/src/gpu/transpose.cu index 1ccba1a0..70d49b82 100644 --- a/src/gpu/transpose.cu +++ b/src/gpu/transpose.cu @@ -255,6 +255,24 @@ transpose_launch(const struct transpose_args& a) return launch(a, first_dim); } +static int +transpose_dispatch(const struct transpose_args& a, uint8_t bpe) +{ + switch (bpe) { + case 1: + return transpose_launch(a); + case 2: + return transpose_launch(a); + case 4: + return transpose_launch(a); + case 8: + return transpose_launch(a); + default: + log_error("transpose: unsupported bpe %u", (unsigned)bpe); + return 1; + } +} + extern "C" int transpose_chunks(CUdeviceptr dst, CUdeviceptr src, @@ -265,7 +283,7 @@ transpose_chunks(CUdeviceptr dst, const struct tile_stream_layout* layout, CUstream stream) { - if (!layout->has_partial_chunks) + if (!layout_has_partial_chunks(layout)) return transpose(dst, src, src_bytes, @@ -289,18 +307,7 @@ transpose_chunks(CUdeviceptr dst, layout->lifted_strides, stream, layout }; - switch (bpe) { - case 1: - return transpose_launch(a); - case 2: - return transpose_launch(a); - case 4: - return transpose_launch(a); - case 8: - return transpose_launch(a); - default: - return 1; - } + return transpose_dispatch(a, bpe); } extern "C" int @@ -320,17 +327,5 @@ transpose(CUdeviceptr d_dst_beg, i_offset, epoch_elements, region_bytes, rank, shape, strides, stream }; - switch (bpe) { - case 1: - return transpose_launch(a); - case 2: - return transpose_launch(a); - case 4: - return transpose_launch(a); - case 8: - return transpose_launch(a); - default: - log_error("transpose: unsupported bpe %u", (unsigned)bpe); - return 1; - } + return transpose_dispatch(a, bpe); } diff --git a/src/stream/config.c b/src/stream/config.c index d26c9f71..8ab9125e 100644 --- a/src/stream/config.c +++ b/src/stream/config.c @@ -90,7 +90,6 @@ compute_level_layout(struct tile_stream_layout* layout, layout->lifted_rank = 2 * rank; layout->chunk_elements = 1; layout->epoch_elements = 1; - layout->has_partial_chunks = 0; uint64_t chunk_count[HALF_MAX_RANK]; for (int i = 0; i < rank; ++i) { @@ -102,7 +101,6 @@ compute_level_layout(struct tile_stream_layout* layout, const uint64_t extent = i < n_append ? dims[i].chunk_size : level_shape[i]; CHECK(Fail, extent > 0); layout->input_shape[i] = extent; - layout->has_partial_chunks |= extent % dims[i].chunk_size != 0; CHECK_MUL_OVERFLOW(Fail, layout->epoch_elements, extent, UINT64_MAX); layout->epoch_elements *= extent; } diff --git a/src/stream/layouts.h b/src/stream/layouts.h index 20af9ab9..516bd737 100644 --- a/src/stream/layouts.h +++ b/src/stream/layouts.h @@ -18,7 +18,6 @@ struct tile_stream_layout // Packed acquisition extents within one epoch. Append dimensions span one // chunk; inner dimensions retain their logical size, including partial edges. uint64_t input_shape[HALF_MAX_RANK]; - int has_partial_chunks; uint64_t chunk_elements; uint64_t chunk_stride; @@ -27,6 +26,14 @@ struct tile_stream_layout size_t chunk_pool_bytes; }; +// Chunk alignment bytes do not affect whether the logical shape fills chunks. +static inline int +layout_has_partial_chunks(const struct tile_stream_layout* layout) +{ + return layout->epoch_elements != + layout->chunks_per_epoch * layout->chunk_elements; +} + // Per-level pre-computed layout information (CPU only, no GPU pointers). struct level_layout_info { diff --git a/tests/test_bench_input.c b/tests/test_bench_input.c index 8d5b0b5b..242632eb 100644 --- a/tests/test_bench_input.c +++ b/tests/test_bench_input.c @@ -77,24 +77,22 @@ test_load_packed(enum dtype dtype) if (fclose(file) || written != source_bytes) goto Cleanup; - if (bench_input_load(&input, path, dtype, 65, 66, 64, 64) || - input.elements != 2 * 65 * 66 || input.frame_elements != 65 * 66 || - input.logical_frame_elements != 65 * 66 || - input.source_bytes != source_bytes) + if (bench_input_load(&input, path, dtype, 65, 66) || + input.frame_elements != 65 * 66 || input.source_bytes != source_bytes) goto Cleanup; if (memcmp(input.data, source, source_bytes)) goto Cleanup; - if (!bench_input_load(&input, path, dtype, SIZE_MAX, 1, 64, 64) || - !bench_input_load(&input, path, dtype, 1, 1, 0, 64) || - !bench_input_load(&input, path, dtype_f64, 65, 66, 64, 64)) - goto Cleanup; bench_input_free(&input); + if (!bench_input_load(&input, path, dtype, SIZE_MAX, 1) || + !bench_input_load(&input, path, dtype, 0, 1) || + !bench_input_load(&input, path, dtype_f64, 65, 66)) + goto Cleanup; file = fopen(path, "wb"); if (!file) goto Cleanup; written = fwrite(source, 1, source_bytes - 1, file); if (fclose(file) || written != source_bytes - 1 || - !bench_input_load(&input, path, dtype, 65, 66, 64, 64) || input.data) + !bench_input_load(&input, path, dtype, 65, 66) || input.data) goto Cleanup; error = 0; From b44050bcf6cc4c8b3441c2e5810d9b75ded00947 Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Sun, 4 Oct 2026 11:08:40 -0700 Subject: [PATCH 3/4] Use shared lookup scatter for packed chunk edges --- README.md | 4 + src/cpu/lod.cpp | 96 +++++++++++----- src/cpu/lod.h | 13 +++ src/cpu/stream.body.c | 42 +++++-- src/cpu/stream.body.h | 1 + src/cpu/stream.c | 17 ++- src/cpu/stream.internal.h | 2 + src/cpu/transpose.cpp | 66 ++--------- src/cpu/transpose.h | 4 +- src/gpu/lod.cu | 139 +++++++++++++++-------- src/gpu/lod.h | 15 +++ src/gpu/stream.c | 1 + src/gpu/stream.engine.h | 2 + src/gpu/stream.ingest.c | 38 +++++-- src/gpu/stream.ingest.h | 1 + src/gpu/stream.init.c | 2 +- src/gpu/stream.lod.c | 24 +++- src/gpu/stream.lod.h | 4 +- src/gpu/transpose.cu | 124 ++++----------------- src/gpu/transpose.h | 13 --- src/multiarray/stream.c | 10 ++ src/stream.cpu.h | 3 +- src/stream.gpu.h | 2 +- src/stream/config.c | 39 +++++++ src/stream/layouts.h | 17 +++ src/util/index.ops.h | 13 +++ tests/CMakeLists.txt | 7 +- tests/test_packed_input.c | 120 ++++++++++++++++++++ tests/test_scatter_lut.c | 225 ++++++++++++++++++++++++++++++++++++++ 29 files changed, 765 insertions(+), 279 deletions(-) create mode 100644 tests/test_scatter_lut.c diff --git a/README.md b/README.md index 3b141cb4..13718376 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,10 @@ LOD pyramids built on the fly: after the base level (L0) is chunked, the pipelin scatters, reduces, and chunks each coarser level before compressing and delivering it alongside L0. +Input remains tightly packed. At partial chunk edges, precomputed row and +column offsets map logical coordinates into zeroed output chunks through the +same lookup scatter used for LOD. Padding requires no work from the caller. + **Supported element types:** u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64 (see `enum dtype` in `src/dtype.h`). diff --git a/src/cpu/lod.cpp b/src/cpu/lod.cpp index dc923367..02f16066 100644 --- a/src/cpu/lod.cpp +++ b/src/cpu/lod.cpp @@ -387,23 +387,64 @@ reduce_typed(const lod_plan* p, } } -template +template static void -morton_to_chunks_typed(const T* values, - T* chunks, - const uint32_t* chunk_lut, - const uint64_t* batch_offsets, - uint64_t lod_count, - uint64_t batch_count, - struct threadpool* pool) -{ - // Same flattening as gather_typed. - uint64_t total = batch_count * lod_count; - run_for_n(pool, total, [=](uint64_t k) { - uint64_t b = k / lod_count; - uint64_t i = k % lod_count; - chunks[batch_offsets[b] + chunk_lut[i]] = values[b * lod_count + i]; - }); +scatter_lut_typed(const T* values, + T* chunks, + const Index* chunk_lut, + const uint64_t* batch_offsets, + uint64_t width, + uint64_t count, + uint64_t i_offset, + struct threadpool* pool) +{ + // Decompose once per row, including ranges split by the thread pool. + auto body = [=](size_t beg, size_t end) { + while (beg < end) { + const uint64_t index = i_offset + beg; + const uint64_t column = index % width; + const uint64_t remaining = width - column; + const uint64_t n = remaining < end - beg ? remaining : end - beg; + T* dst = chunks + batch_offsets[index / width]; + for (uint64_t i = 0; i < n; ++i) + dst[chunk_lut[column + i]] = values ? values[beg + i] : T{}; + beg += n; + } + }; + auto trampoline = [](size_t beg, size_t end, int tid, void* ctx) { + (void)tid; + (*static_cast(ctx))(beg, end); + }; + if (count * sizeof(T) < (64u << 10)) + body(0, count); + else + threadpool_for_n(pool, count, trampoline, &body); +} + +extern "C" int +scatter_lut_cpu(void* dst, + const void* src, + uint64_t count, + uint8_t bpe, + uint64_t i_offset, + uint64_t width, + const uint64_t* columns, + const uint64_t* rows, + struct threadpool* pool) +{ +#define CASE(B, T) \ + case B: \ + scatter_lut_typed( \ + (const T*)src, (T*)dst, columns, rows, width, count, i_offset, pool); \ + return 0 + switch (bpe) { + CASE(1, uint8_t); + CASE(2, uint16_t); + CASE(4, uint32_t); + CASE(8, uint64_t); + } +#undef CASE + return 1; } // ---- Morton-to-chunks LUT ---- @@ -431,10 +472,10 @@ build_chunk_lut(const lod_plan* p, uint64_t coord = coords[d]; int full_d = ld->lod_to_dim[d]; uint64_t chunk_size_d = layout->lifted_shape[2 * full_d + 1]; - uint64_t chunk_idx = coord / chunk_size_d; - uint64_t within = coord % chunk_size_d; - offset += (int64_t)chunk_idx * layout->lifted_strides[2 * full_d]; - offset += (int64_t)within * layout->lifted_strides[2 * full_d + 1]; + offset += chunk_coordinate_offset(coord, + chunk_size_d, + layout->lifted_strides[2 * full_d], + layout->lifted_strides[2 * full_d + 1]); } return (uint32_t)offset; @@ -604,13 +645,14 @@ lod_cpu_morton_to_chunks(const lod_plan* p, } #define DO(T) \ - morton_to_chunks_typed((const T*)lv_values, \ - (T*)chunk_pool, \ - chunk_lut, \ - fixed_dims_chunk_offsets, \ - lod_count, \ - p->levels.level[lv].fixed_dims_count, \ - pool) + scatter_lut_typed((const T*)lv_values, \ + (T*)chunk_pool, \ + chunk_lut, \ + fixed_dims_chunk_offsets, \ + lod_count, \ + p->levels.level[lv].fixed_dims_count* lod_count, \ + 0, \ + pool) DISPATCH(dtype, DO); #undef DO diff --git a/src/cpu/lod.h b/src/cpu/lod.h index 14c9a7bc..f6a96f22 100644 --- a/src/cpu/lod.h +++ b/src/cpu/lod.h @@ -16,6 +16,19 @@ extern "C" struct threadpool; + // Scatter a source range through a row-plus-column lookup. i_offset is + // within one epoch; source starts at the supplied range. NULL src clears + // the mapped elements (used for the final partial epoch). + int scatter_lut_cpu(void* dst, + const void* src, + uint64_t count, + uint8_t bpe, + uint64_t i_offset, + uint64_t width, + const uint64_t* columns, + const uint64_t* rows, + struct threadpool* pool); + // Reduce across LOD levels in-place. // values buffer holds all levels: total = levels.ends[nlod-1] elements. // csrs: array of nlod-1 reduce_csr entries, one per level transition. diff --git a/src/cpu/stream.body.c b/src/cpu/stream.body.c index 43a231ac..4ace6243 100644 --- a/src/cpu/stream.body.c +++ b/src/cpu/stream.body.c @@ -1,5 +1,6 @@ #include "cpu/stream.body.h" +#include "cpu/lod.h" #include "cpu/transpose.h" #include "platform/platform.h" #include "util/metric.h" @@ -90,6 +91,30 @@ validate_view(const struct cpu_stream_view* v) #define validate_view(v) ((void)0) #endif +static int +scatter_input(struct cpu_stream_view* v, + void* dst, + const void* src, + uint64_t bytes, + uint64_t i_offset) +{ + const uint8_t bpe = (uint8_t)dtype_bpe(v->config->dtype); + if (v->input_chunk_lut) { + const uint64_t width = + v->layout->input_shape[v->layout->lifted_rank / 2 - 1]; + return scatter_lut_cpu(dst, + src, + bytes / bpe, + bpe, + i_offset % v->layout->epoch_elements, + width, + v->input_chunk_lut, + v->input_chunk_lut + width, + v->pool); + } + return transpose_cpu(dst, src, bytes, bpe, i_offset, v->layout, v->pool); +} + // ---- Shared append body ---- struct writer_result @@ -139,13 +164,8 @@ cpu_stream_append_body(struct cpu_stream_view* v, struct slice input) v->levels->total_chunks * v->layout->chunk_stride * bpe; CHECK(Error, - transpose_cpu(epoch_pool, - src, - bytes, - (uint8_t)bpe, - *v->cursor_elements, - v->layout, - v->pool) == 0); + scatter_input(v, epoch_pool, src, bytes, *v->cursor_elements) == + 0); } float ms = (float)(platform_toc(&clk) * 1000.0); @@ -269,13 +289,11 @@ cpu_stream_flush_body(struct cpu_stream_view* v) v->levels->total_chunks * v->layout->chunk_stride * bpe; CHECK(Fail, - transpose_cpu(epoch_pool, + scatter_input(v, + epoch_pool, NULL, (v->layout->epoch_elements - used) * bpe, - (uint8_t)bpe, - used, - v->layout, - v->pool) == 0); + used) == 0); } v->batch_active_masks[*v->batch_accumulated] = active_mask; (*v->batch_accumulated)++; diff --git a/src/cpu/stream.body.h b/src/cpu/stream.body.h index a1598231..1c7be041 100644 --- a/src/cpu/stream.body.h +++ b/src/cpu/stream.body.h @@ -31,6 +31,7 @@ struct cpu_stream_view struct shard_state* shard; // [LOD_MAX_LEVELS] array struct aggregate_layout* agg_layout; // [LOD_MAX_LEVELS] array struct reduce_csr* csrs; // [nlod-1] CSR LUTs + const uint64_t* input_chunk_lut; // direct row/column map, or NULL void* append_accum; uint32_t* append_counts; // [LOD_MAX_LEVELS] diff --git a/src/cpu/stream.c b/src/cpu/stream.c index 15e0f028..94857670 100644 --- a/src/cpu/stream.c +++ b/src/cpu/stream.c @@ -64,6 +64,14 @@ tile_stream_cpu_create(const struct tile_stream_configuration* config, s->layout = s->cl.layouts[0]; s->levels = s->cl.levels; + if (!s->levels.enable_multiscale && layout_has_partial_chunks(&s->layout)) { + const size_t bytes = chunk_scatter_lut_bytes(&s->layout); + CHECK(Fail, bytes); + s->input_chunk_lut = (uint64_t*)malloc(bytes); + CHECK(Fail, s->input_chunk_lut); + chunk_scatter_lut_build(&s->layout, s->input_chunk_lut); + } + const uint32_t K = s->cl.epochs_per_batch; const size_t bytes_per_element = dtype_bpe(config->dtype); const uint64_t total_chunks = s->levels.total_chunks; @@ -350,6 +358,7 @@ tile_stream_cpu_release_resources(struct tile_stream_cpu* s) free(s->comp_sizes); free(s->batch_active_masks); free(s->pool_epochs_scratch); + free(s->input_chunk_lut); free(s->scatter_lut); free(s->scatter_fixed_dims_offsets); free(s->linear); @@ -519,9 +528,14 @@ compute_memory_info(const struct computed_stream_layouts* cl, info->aggregate_bytes = agg; } - // LOD buffers (multiscale only). + // Scatter LUTs and multiscale buffers. { size_t lod = 0; + if (!cl->levels.enable_multiscale && + layout_has_partial_chunks(&cl->layouts[0])) { + lod = chunk_scatter_lut_bytes(&cl->layouts[0]); + CHECK(Error, lod); + } if (cl->levels.enable_multiscale) { lod += cl->layouts[0].epoch_elements * bytes_per_element; // linear uint64_t total_lod_elements = @@ -793,6 +807,7 @@ make_view(struct tile_stream_cpu* s) .batch_active_masks = s->batch_active_masks, .pool_epochs_scratch = s->pool_epochs_scratch, .pool_fully_covered = s->pool_fully_covered, + .input_chunk_lut = s->input_chunk_lut, .shard = s->shard, .agg_layout = s->agg_layout, .csrs = s->csrs, diff --git a/src/cpu/stream.internal.h b/src/cpu/stream.internal.h index 38bb5a13..6a868b6f 100644 --- a/src/cpu/stream.internal.h +++ b/src/cpu/stream.internal.h @@ -20,6 +20,8 @@ struct tile_stream_cpu struct tile_stream_layout layout; struct level_geometry levels; + uint64_t* input_chunk_lut; // direct row/column map (single-level edges) + // Chunk pool: total_chunks * chunk_stride * bpe bytes. void* chunk_pool; diff --git a/src/cpu/transpose.cpp b/src/cpu/transpose.cpp index 39eba67b..8627bf6b 100644 --- a/src/cpu/transpose.cpp +++ b/src/cpu/transpose.cpp @@ -21,7 +21,7 @@ scatter_loop(T* dst, { const T* s = (const T*)my_src; for (uint64_t i = 0; i < my_n; ++i) { - dst[o] = s[i]; + dst[o] = s ? s[i] : T{}; o += inner_stride; if (++coords[rank - 1] >= shape[rank - 1]) { coords[rank - 1] = 0; @@ -47,57 +47,8 @@ struct transpose_ctx int64_t inner_stride; uint64_t epoch_elements; uint8_t bpe; - const struct tile_stream_layout* layout; }; -// Decode logical coordinates once per run ending at a row or chunk edge. -// The destination pool is already zeroed; cells outside input_shape are never -// visited. No padded frame or per-element division is needed. -static void -transpose_packed_range(size_t beg, size_t end, int tid, void* vctx) -{ - (void)tid; - const auto* c = (const transpose_ctx*)vctx; - const auto* l = c->layout; - const int rank = l->lifted_rank / 2; - while (beg < end) { - uint64_t rest = (c->i_offset + beg) % l->epoch_elements; - uint64_t count = end - beg; - int64_t out = 0; - for (int d = rank - 1; d >= 0; --d) { - const uint64_t coord = rest % l->input_shape[d]; - rest /= l->input_shape[d]; - const uint64_t chunk = l->lifted_shape[2 * d + 1]; - out += (coord / chunk) * l->lifted_strides[2 * d] + - (coord % chunk) * l->lifted_strides[2 * d + 1]; - if (d == rank - 1) { - const uint64_t row_left = l->input_shape[d] - coord; - const uint64_t chunk_left = chunk - coord % chunk; - if (count > row_left) - count = row_left; - if (count > chunk_left) - count = chunk_left; - } - } - char* dst = (char*)c->dst + out * c->bpe; - const char* src = c->src ? c->src + beg * c->bpe : nullptr; - const int64_t step = l->lifted_strides[l->lifted_rank - 1] * c->bpe; - if (step == c->bpe) { - if (src) - memcpy(dst, src, count * c->bpe); - else - memset(dst, 0, count * c->bpe); - } else - for (uint64_t i = 0; i < count; ++i) { - if (src) - memcpy(dst + i * step, src + i * c->bpe, c->bpe); - else - memset(dst + i * step, 0, c->bpe); - } - beg += count; - } -} - // The append chunk coordinates have zero strides: the caller selects their // destination epoch. All remaining coordinates describe one epoch. If their // strides are row-major, every range within the epoch can be copied directly. @@ -130,10 +81,13 @@ transpose_range(size_t beg, size_t end, int tid, void* vctx) uint64_t coords[MAX_RANK]; int64_t o = (int64_t)transposed_offset(c->rank, c->shape, c->strides, base, coords); - const void* my_src = c->src + beg * c->bpe; + const void* my_src = c->src ? c->src + beg * c->bpe : nullptr; if (range_is_contiguous(c, base, my_n)) { - memcpy((char*)c->dst + o * c->bpe, my_src, my_n * c->bpe); + if (my_src) + memcpy((char*)c->dst + o * c->bpe, my_src, my_n * c->bpe); + else + memset((char*)c->dst + o * c->bpe, 0, my_n * c->bpe); return; } @@ -194,18 +148,14 @@ transpose_cpu(void* dst, inner_stride, layout->epoch_elements, bpe, - layout, }; // For small appends, dispatching and joining the pool costs more than the // scatter. Keep this choice independent of compression parallelism. constexpr uint64_t min_parallel_bytes = 64u << 10; - auto range = layout_has_partial_chunks(layout) || !src - ? transpose_packed_range - : transpose_range; if (src_bytes < min_parallel_bytes) - range(0, n, 0, &c); + transpose_range(0, n, 0, &c); else - threadpool_for_n(pool, n, range, &c); + threadpool_for_n(pool, n, transpose_range, &c); return 0; } diff --git a/src/cpu/transpose.h b/src/cpu/transpose.h index f9cc34f5..d08f2561 100644 --- a/src/cpu/transpose.h +++ b/src/cpu/transpose.h @@ -11,9 +11,9 @@ extern "C" struct threadpool; struct tile_stream_layout; - // CPU scatter transpose using the vadd() algorithm. + // CPU scatter transpose for chunk-aligned input using the vadd() algorithm. // Copies directly when the layout proves that the epoch is contiguous; - // otherwise maps packed logical coordinates into padded output chunks. + // otherwise walks its lifted shape and strides. // i_offset is the global flat input offset (for multi-call accumulation). // A NULL src zeroes the mapped range (used for a final partial epoch). // Returns 0 on success and nonzero for an unsupported element size. diff --git a/src/gpu/lod.cu b/src/gpu/lod.cu index 957f0d3d..01839e3c 100644 --- a/src/gpu/lod.cu +++ b/src/gpu/lod.cu @@ -384,10 +384,10 @@ lod_build_chunk_scatter_lut_k(uint32_t* __restrict__ chunk_lut, uint64_t coord = remainder % lod_shape[d]; remainder /= lod_shape[d]; coords[d] = coord; - uint64_t chunk_idx = coord / lod_chunk_sizes[d]; - uint64_t within = coord % lod_chunk_sizes[d]; - offset += (int64_t)chunk_idx * lod_chunk_strides[2 * d]; - offset += (int64_t)within * lod_chunk_strides[2 * d + 1]; + offset += chunk_coordinate_offset(coord, + lod_chunk_sizes[d], + lod_chunk_strides[2 * d], + lod_chunk_strides[2 * d + 1]); } uint64_t morton_pos = @@ -395,27 +395,48 @@ lod_build_chunk_scatter_lut_k(uint32_t* __restrict__ chunk_lut, chunk_lut[morton_pos] = (uint32_t)offset; } -// --- Morton-to-chunk scatter kernel using LUT --- -// Sequential reads from morton buffer, LUT-directed writes to chunk pool. +// --- Chunk scatter using LUTs --- +// LOD and direct input use the same row-plus-column lookup. Repetition adds +// the destination epoch base without changing how a table maps its elements. -template +template +struct scatter_lut_view +{ + const Index* columns; + const Index* rows; + uint64_t width; + + __device__ uint64_t operator()(uint64_t i) const + { + return (uint64_t)rows[i / width] + columns[i % width]; + } +}; + +struct repeated_scatter_lut_view +{ + scatter_lut_view lut; + uint64_t epoch_elements; + uint64_t region_stride; + + __device__ uint64_t operator()(uint64_t i) const + { + return (i / epoch_elements) * region_stride + lut(i % epoch_elements); + } +}; + +template __global__ void -lod_morton_to_chunks_lut_k( - T* __restrict__ dst, - const T* __restrict__ src, - const uint32_t* __restrict__ chunk_lut, - const uint32_t* __restrict__ fixed_dims_chunk_offsets, - uint64_t lod_count, - uint64_t total) +scatter_lut_k(T* __restrict__ dst, + const T* __restrict__ src, + Mapping mapping, + uint64_t total, + uint64_t i_offset) { const uint64_t gid = (uint64_t)blockIdx.x * LOD_BLOCK + threadIdx.x; if (gid >= total) return; - uint64_t batch = gid / lod_count; - uint64_t morton_pos = gid % lod_count; - dst[(uint64_t)fixed_dims_chunk_offsets[batch] + chunk_lut[morton_pos]] = - src[gid]; + dst[mapping(i_offset + gid)] = src[gid]; } // --- Shared memory size helpers --- @@ -528,27 +549,24 @@ lod_build_chunk_scatter_lut(CUdeviceptr d_chunk_lut, return 1; } -template +template static int -lod_morton_to_chunks_lut_launch(CUdeviceptr d_chunks, - CUdeviceptr d_morton, - CUdeviceptr d_chunk_lut, - CUdeviceptr d_fixed_dims_chunk_offsets, - uint64_t lod_count, - uint64_t fixed_dims_count, - CUstream stream) +scatter_lut_launch(CUdeviceptr dst, + CUdeviceptr src, + Mapping mapping, + uint64_t count, + uint64_t i_offset, + CUstream stream) { - const uint64_t total = fixed_dims_count * lod_count; - const int grid_size = (int)((total + LOD_BLOCK - 1) / LOD_BLOCK); - - return CUDA_LAUNCH(lod_morton_to_chunks_lut_k - <<>>( - (T*)d_chunks, - (const T*)d_morton, - (const uint32_t*)d_chunk_lut, - (const uint32_t*)d_fixed_dims_chunk_offsets, - lod_count, - total)); + if (count == 0) + return 0; + const uint64_t blocks = ceildiv(count, (uint64_t)LOD_BLOCK); + if (blocks > INT32_MAX || i_offset > UINT64_MAX - (count - 1)) + return 1; + const unsigned grid_size = (unsigned)blocks; + + return CUDA_LAUNCH(scatter_lut_k<<>>( + (T*)dst, (const T*)src, mapping, count, i_offset)); } extern "C" int @@ -561,21 +579,54 @@ lod_morton_to_chunks_lut(CUdeviceptr d_chunks, uint64_t fixed_dims_count, CUstream stream) { + const scatter_lut_view mapping = { + (const uint32_t*)d_chunk_lut, + (const uint32_t*)d_fixed_dims_chunk_offsets, + lod_count, + }; #define DISPATCH(D, T) \ if (dtype == D) { \ - return lod_morton_to_chunks_lut_launch(d_chunks, \ - d_morton, \ - d_chunk_lut, \ - d_fixed_dims_chunk_offsets, \ - lod_count, \ - fixed_dims_count, \ - stream); \ + return scatter_lut_launch( \ + d_chunks, d_morton, mapping, fixed_dims_count * lod_count, 0, stream); \ } FOR_EACH_DTYPE(DISPATCH) #undef DISPATCH return 1; } +extern "C" int +scatter_lut_gpu(CUdeviceptr dst, + CUdeviceptr src, + uint64_t count, + uint8_t bpe, + uint64_t i_offset, + uint64_t width, + uint64_t epoch_elements, + uint64_t region_bytes, + CUdeviceptr columns, + CUdeviceptr rows, + CUstream stream) +{ + if (!width || !epoch_elements || !bpe || region_bytes % bpe || src % bpe) + return 1; + const repeated_scatter_lut_view mapping = { + { (const uint64_t*)columns, (const uint64_t*)rows, width }, + epoch_elements, + region_bytes / bpe, + }; +#define CASE(B, T) \ + case B: \ + return scatter_lut_launch(dst, src, mapping, count, i_offset, stream) + switch (bpe) { + CASE(1, uint8_t); + CASE(2, uint16_t); + CASE(4, uint32_t); + CASE(8, uint64_t); + } +#undef CASE + return 1; +} + template static int lod_gather_lut_launch(CUdeviceptr d_dst, diff --git a/src/gpu/lod.h b/src/gpu/lod.h index 2656af4f..a26d459a 100644 --- a/src/gpu/lod.h +++ b/src/gpu/lod.h @@ -10,6 +10,21 @@ extern "C" { #endif + // Direct row-plus-column lookup scatter. The tables repeat each epoch; + // region_bytes is the distance between output epochs. Offsets are uint64_t + // elements, and input starts at i_offset within the first epoch. + int scatter_lut_gpu(CUdeviceptr dst, + CUdeviceptr src, + uint64_t count, + uint8_t bpe, + uint64_t i_offset, + uint64_t width, + uint64_t epoch_elements, + uint64_t region_bytes, + CUdeviceptr columns, + CUdeviceptr rows, + CUstream stream); + // CSR-based reduce: precomputed starts/indices LUT. // d_values is a single flat allocation; src_offset/dst_offset are element // counts into it. diff --git a/src/gpu/stream.c b/src/gpu/stream.c index a73369ec..2e3d2a90 100644 --- a/src/gpu/stream.c +++ b/src/gpu/stream.c @@ -117,6 +117,7 @@ engine_dispatch_ingest(struct stream_engine* e, .first_epoch = stream_engine_pool_epoch(e, ctx, e->sched.accumulated), .epoch_bytes = pool_epoch_bytes(ctx), .epochs = e->sched.epochs_per_batch - e->sched.accumulated, + .input_chunk_lut = e->lod.d_input_chunk_lut, }; return ingest_dispatch_scatter(&e->stage, &ctx->layout, diff --git a/src/gpu/stream.engine.h b/src/gpu/stream.engine.h index f5a9d184..43531885 100644 --- a/src/gpu/stream.engine.h +++ b/src/gpu/stream.engine.h @@ -114,6 +114,8 @@ struct lod_state { struct lod_plan plan; + CUdeviceptr d_input_chunk_lut; // u64 row/column map (single-level edges) + CUdeviceptr d_full_shape; // device copy of shapes[0] CUdeviceptr d_lod_shape; // device copy of LOD-projected shapes[0] CUdeviceptr d_gather_lut; // u32, lod_nelem[0] entries diff --git a/src/gpu/stream.ingest.c b/src/gpu/stream.ingest.c index 4e418216..2b8e3280 100644 --- a/src/gpu/stream.ingest.c +++ b/src/gpu/stream.ingest.c @@ -1,6 +1,7 @@ #include "gpu/stream.ingest.h" #include "gpu/host_memory.h" +#include "gpu/lod.h" #include "gpu/metric.cuda.h" #include "gpu/prelude.cuda.h" #include "gpu/transpose.h" @@ -169,15 +170,34 @@ scatter_with_timing(struct staging_state* stage, { struct scatter_timing* st = timing_begin(stage, bytes, compute); CHECK_SILENT(Fail, st); - CHECK_SILENT(Fail, - transpose_chunks(gpu_pool_view_d(dst.first_epoch), - gpu_pool_view_d(d_in), - bytes, - (uint8_t)bpe, - in_epoch, - dst.epoch_bytes, - layout, - compute) == 0); + if (dst.input_chunk_lut) { + const uint64_t width = layout->input_shape[layout->lifted_rank / 2 - 1]; + CHECK_SILENT(Fail, + scatter_lut_gpu(gpu_pool_view_d(dst.first_epoch), + gpu_pool_view_d(d_in), + bytes / bpe, + (uint8_t)bpe, + in_epoch, + width, + layout->epoch_elements, + dst.epoch_bytes, + dst.input_chunk_lut, + dst.input_chunk_lut + width * sizeof(uint64_t), + compute) == 0); + } else { + CHECK_SILENT(Fail, + transpose(gpu_pool_view_d(dst.first_epoch), + gpu_pool_view_d(d_in), + bytes, + (uint8_t)bpe, + in_epoch, + layout->epoch_elements, + dst.epoch_bytes, + layout->lifted_rank, + layout->lifted_shape, + layout->lifted_strides, + compute) == 0); + } return timing_end(st, compute); Fail: diff --git a/src/gpu/stream.ingest.h b/src/gpu/stream.ingest.h index e4347dcc..0fa6d73f 100644 --- a/src/gpu/stream.ingest.h +++ b/src/gpu/stream.ingest.h @@ -46,6 +46,7 @@ struct scatter_destination // the caller acquired size_t epoch_bytes; // one epoch's region in the pool uint32_t epochs; // regions the caller has, counting from first_epoch + CUdeviceptr input_chunk_lut; // u64 row/column map, or 0 for lifted transpose }; // H2D transfer + scatter into chunk pool. diff --git a/src/gpu/stream.init.c b/src/gpu/stream.init.c index 61920450..6752de86 100644 --- a/src/gpu/stream.init.c +++ b/src/gpu/stream.init.c @@ -244,7 +244,7 @@ engine_array_state_init(struct engine_array_state* st, st->lod.layouts[lv] = cl->layouts[lv]; // The lifted fast path needs trailing extents spanning exactly one epoch. - // Partial chunks use explicit logical coordinates in transpose_chunks. + // Partial chunks use precomputed logical-to-chunk lookup tables. for (int lv = 0; lv < cl->levels.nlod; ++lv) { const struct tile_stream_layout* l = &st->lod.layouts[lv]; if (!layout_has_partial_chunks(l) && diff --git a/src/gpu/stream.lod.c b/src/gpu/stream.lod.c index 8cfcabbe..fed86e57 100644 --- a/src/gpu/stream.lod.c +++ b/src/gpu/stream.lod.c @@ -338,8 +338,25 @@ lod_state_init(struct lod_state* lod, struct level_geometry* levels, const struct tile_stream_configuration* config) { - if (!levels->enable_multiscale) + if (!levels->enable_multiscale) { + const struct tile_stream_layout* layout = &lod->layouts[0]; + if (layout_has_partial_chunks(layout)) { + const size_t bytes = chunk_scatter_lut_bytes(layout); + CHECK(Fail, bytes); + uint64_t* lut = (uint64_t*)malloc(bytes); + CHECK(Fail, lut); + chunk_scatter_lut_build(layout, lut); + CUresult rc = cuMemAlloc(&lod->d_input_chunk_lut, bytes); + if (rc == CUDA_SUCCESS) + rc = cuMemcpyHtoD(lod->d_input_chunk_lut, lut, bytes); + // Publish setup before nonblocking compute streams can read the table. + if (rc == CUDA_SUCCESS) + rc = cuStreamSynchronize(0); + free(lut); + CU(Fail, rc); + } return 0; + } // LOD-specific: plan shapes, gather LUT, reduce arrays, morton LUTs. for (int k = 0; k < lod->plan.levels.nlod; ++k) { @@ -499,7 +516,9 @@ lod_state_device_bytes(const struct computed_stream_layouts* cl, const struct lod_plan* p = &cl->plan; if (!cl->levels.enable_multiscale) - return 0; + return layout_has_partial_chunks(&cl->layouts[0]) + ? chunk_scatter_lut_bytes(&cl->layouts[0]) + : 0; size_t bytes = 0; @@ -567,6 +586,7 @@ lod_state_destroy(struct lod_state* lod) CUWARN(cuMemFree(lod->append_accum.d_counts)); // Per-array LOD allocations + CUWARN(cuMemFree(lod->d_input_chunk_lut)); CUWARN(cuMemFree(lod->d_full_shape)); CUWARN(cuMemFree(lod->d_lod_shape)); CUWARN(cuMemFree(lod->d_gather_lut)); diff --git a/src/gpu/stream.lod.h b/src/gpu/stream.lod.h index 6bbe1751..45b517bf 100644 --- a/src/gpu/stream.lod.h +++ b/src/gpu/stream.lod.h @@ -3,8 +3,8 @@ #include "gpu/stream.internal.h" #include "stream/dim_info.h" -// When multiscale is enabled, uploads LOD plan shapes and builds -// scatter/reduce LUTs; does nothing otherwise. +// Builds a direct scatter LUT for single-level edges, or uploads LOD plan +// shapes and builds scatter/reduce LUTs when multiscale is enabled. // Plan and level layouts must already be populated in lod->plan and // lod->layouts (from compute_stream_layouts). Sets levels->nlod. // Returns 0 on success. diff --git a/src/gpu/transpose.cu b/src/gpu/transpose.cu index 70d49b82..bbb0a9ac 100644 --- a/src/gpu/transpose.cu +++ b/src/gpu/transpose.cu @@ -1,7 +1,6 @@ #include "defs.limits.h" #include "gpu/prelude.cuda.h" #include "gpu/transpose.h" -#include "stream/layouts.h" #include "util/prelude.h" #include @@ -11,19 +10,8 @@ constexpr int ELEMENTS_PER_BLOCK = (1 << 12) / (int)sizeof(T); // Travels in the launch parameters, so the kernel reads it from the constant // bank instead of paying a load per dimension per element. -template -struct scatter_edges -{}; - template -struct scatter_edges -{ - Index chunks[HALF_MAX_RANK]; - Index chunk_strides[HALF_MAX_RANK]; -}; - -template -struct scatter_layout : scatter_edges +struct scatter_layout { Index shape[MAX_RANK]; Index strides[MAX_RANK]; @@ -31,14 +19,14 @@ struct scatter_layout : scatter_edges }; // Transpose data kernel - v0 -template +template __global__ void __launch_bounds__(256, 4) transpose_v0_k(T* d_dst, const T* d_src, Index src_size, Index i_offset, Index region_stride, - scatter_layout layout) + scatter_layout layout) { const Index block_offset = (Index)blockIdx.x * ELEMENTS_PER_BLOCK; const Index left = src_size - block_offset; @@ -52,11 +40,7 @@ __global__ void __launch_bounds__(256, 4) for (int d = layout.ndims - 1; d >= 0; --d) { const Index coord = rest % layout.shape[d]; rest /= layout.shape[d]; - if constexpr (Packed) - out_offset += (coord / layout.chunks[d]) * layout.chunk_strides[d] + - (coord % layout.chunks[d]) * layout.strides[d]; - else - out_offset += coord * layout.strides[d]; + out_offset += coord * layout.strides[d]; } // The visited extents span one epoch, so what is left counts epochs. @@ -76,7 +60,6 @@ struct transpose_args const uint64_t* shape; const int64_t* strides; CUstream stream; - const struct tile_stream_layout* packed = nullptr; }; template @@ -175,17 +158,6 @@ indices_fit_32_bits(const struct transpose_args& a, if (!add_within_32_bits(&max_offset, last / a.epoch_elements, region_stride)) return 0; - if (a.packed) { - for (int d = 0; d < a.rank; ++d) - if (a.shape[d] > UINT32_MAX || - a.packed->lifted_shape[2 * d + 1] > UINT32_MAX || - (uint64_t)a.strides[2 * d] > UINT32_MAX || - (uint64_t)a.strides[2 * d + 1] > UINT32_MAX) - return 0; - return add_within_32_bits( - &max_offset, 1, a.packed->chunks_per_epoch * a.packed->chunk_stride - 1); - } - for (int d = first_dim; d < a.rank; ++d) { if (a.shape[d] > UINT32_MAX || a.strides[d] < 0) return 0; @@ -196,30 +168,25 @@ indices_fit_32_bits(const struct transpose_args& a, return 1; } -template +template static int launch(const struct transpose_args& a, int first_dim) { const uint64_t src_size = a.src_bytes / sizeof(T); const uint64_t region_stride = a.region_bytes / sizeof(T); - scatter_layout layout = {}; + scatter_layout layout = {}; layout.ndims = a.rank - first_dim; for (int d = 0; d < layout.ndims; ++d) { layout.shape[d] = (Index)a.shape[first_dim + d]; - if constexpr (Packed) { - layout.strides[d] = (Index)a.strides[2 * d + 1]; - layout.chunks[d] = (Index)a.packed->lifted_shape[2 * d + 1]; - layout.chunk_strides[d] = (Index)a.strides[2 * d]; - } else - layout.strides[d] = (Index)a.strides[first_dim + d]; + layout.strides[d] = (Index)a.strides[first_dim + d]; } const int block_size = 256; const unsigned grid_size = (unsigned)ceildiv(src_size, (uint64_t)ELEMENTS_PER_BLOCK); - return CUDA_LAUNCH(transpose_v0_k + return CUDA_LAUNCH(transpose_v0_k <<>>( (T*)a.d_dst_beg, (const T*)a.d_src_beg, @@ -239,12 +206,6 @@ transpose_launch(const struct transpose_args& a) if (!args_valid(a)) return 1; - if (a.packed) { - if (indices_fit_32_bits(a, src_size, 0, a.region_bytes / sizeof(T))) - return launch(a, 0); - return launch(a, 0); - } - const int first_dim = first_decomposed_dim(a.epoch_elements, a.rank, a.shape, a.strides); if (first_dim < 0) @@ -255,61 +216,6 @@ transpose_launch(const struct transpose_args& a) return launch(a, first_dim); } -static int -transpose_dispatch(const struct transpose_args& a, uint8_t bpe) -{ - switch (bpe) { - case 1: - return transpose_launch(a); - case 2: - return transpose_launch(a); - case 4: - return transpose_launch(a); - case 8: - return transpose_launch(a); - default: - log_error("transpose: unsupported bpe %u", (unsigned)bpe); - return 1; - } -} - -extern "C" int -transpose_chunks(CUdeviceptr dst, - CUdeviceptr src, - uint64_t src_bytes, - uint8_t bpe, - uint64_t i_offset, - uint64_t region_bytes, - const struct tile_stream_layout* layout, - CUstream stream) -{ - if (!layout_has_partial_chunks(layout)) - return transpose(dst, - src, - src_bytes, - bpe, - i_offset, - layout->epoch_elements, - region_bytes, - layout->lifted_rank, - layout->lifted_shape, - layout->lifted_strides, - stream); - - const struct transpose_args a = { dst, - src, - src_bytes, - i_offset, - layout->epoch_elements, - region_bytes, - (uint8_t)(layout->lifted_rank / 2), - layout->input_shape, - layout->lifted_strides, - stream, - layout }; - return transpose_dispatch(a, bpe); -} - extern "C" int transpose(CUdeviceptr d_dst_beg, CUdeviceptr d_src_beg, @@ -327,5 +233,17 @@ transpose(CUdeviceptr d_dst_beg, i_offset, epoch_elements, region_bytes, rank, shape, strides, stream }; - return transpose_dispatch(a, bpe); + switch (bpe) { + case 1: + return transpose_launch(a); + case 2: + return transpose_launch(a); + case 4: + return transpose_launch(a); + case 8: + return transpose_launch(a); + default: + log_error("transpose: unsupported bpe %u", (unsigned)bpe); + return 1; + } } diff --git a/src/gpu/transpose.h b/src/gpu/transpose.h index 3d40b05d..0b7f8b9c 100644 --- a/src/gpu/transpose.h +++ b/src/gpu/transpose.h @@ -10,19 +10,6 @@ extern "C" { #endif - struct tile_stream_layout; - - // Map packed acquisition coordinates directly into zeroed chunk pools. - // The range can cross epochs; region_bytes separates destination epochs. - int transpose_chunks(CUdeviceptr dst, - CUdeviceptr src, - uint64_t src_bytes, - uint8_t bpe, - uint64_t i_offset, - uint64_t region_bytes, - const struct tile_stream_layout* layout, - CUstream stream); - // The layout half of what transpose checks, so a stream can refuse a layout // it could never scatter when it opens rather than partway through an // append. Returns 0 when the scatter can place the layout. diff --git a/src/multiarray/stream.c b/src/multiarray/stream.c index a385aaf0..9c84ea80 100644 --- a/src/multiarray/stream.c +++ b/src/multiarray/stream.c @@ -39,6 +39,7 @@ struct array_descriptor struct tile_stream_configuration config; struct computed_stream_layouts cl; struct tile_stream_layout layout; + uint64_t* input_chunk_lut; struct level_geometry levels; struct aggregate_layout agg_layout[LOD_MAX_LEVELS]; struct shard_state shard[LOD_MAX_LEVELS]; @@ -177,6 +178,13 @@ init_array_descriptor(struct array_descriptor* desc, desc->layout = desc->cl.layouts[0]; desc->levels = desc->cl.levels; + if (!desc->levels.enable_multiscale && + layout_has_partial_chunks(&desc->layout)) { + const size_t bytes = chunk_scatter_lut_bytes(&desc->layout); + if (!bytes || !(desc->input_chunk_lut = (uint64_t*)malloc(bytes))) + return 1; + chunk_scatter_lut_build(&desc->layout, desc->input_chunk_lut); + } desc->pool_fully_covered = (desc->layout.chunk_stride == desc->layout.chunk_elements); @@ -517,6 +525,7 @@ multiarray_tile_stream_cpu_release_resources( free(desc->append_accum); free(desc->batch_active_masks); free(desc->pool_epochs_scratch); + free(desc->input_chunk_lut); if (desc->csrs) { int ncsr = desc->cl.plan.levels.nlod - 1; for (int l = 0; l < ncsr; ++l) @@ -670,6 +679,7 @@ make_multiarray_view(struct multiarray_tile_stream_cpu* ms, .batch_active_masks = desc->batch_active_masks, .pool_epochs_scratch = desc->pool_epochs_scratch, .pool_fully_covered = desc->pool_fully_covered, + .input_chunk_lut = desc->input_chunk_lut, .shard = desc->shard, .agg_layout = desc->agg_layout, .csrs = desc->csrs, diff --git a/src/stream.cpu.h b/src/stream.cpu.h index 3130ba81..f464f8a8 100644 --- a/src/stream.cpu.h +++ b/src/stream.cpu.h @@ -50,8 +50,7 @@ struct tile_stream_cpu_memory_info size_t comp_sizes_bytes; // K * total_chunks * sizeof(size_t) size_t aggregate_bytes; // 2x per-batch scratch slots + batch masks size_t host_output_pool_bytes; - size_t lod_bytes; // linear + lod_values + morton_lut + batch_offsets + - // append_accum + size_t lod_bytes; // scatter LUTs and multiscale buffers/accumulators size_t shard_bytes; // active_shard arrays + index buffers uint64_t chunks_per_epoch; diff --git a/src/stream.gpu.h b/src/stream.gpu.h index cd9f90c9..f8417fc9 100644 --- a/src/stream.gpu.h +++ b/src/stream.gpu.h @@ -16,7 +16,7 @@ struct tile_stream_memory_info size_t chunk_pool_bytes; size_t compressed_pool_bytes; size_t aggregate_bytes; - size_t lod_bytes; + size_t lod_bytes; // scatter LUTs and multiscale buffers/accumulators size_t codec_bytes; // codec-owned device allocation bytes // Breakdown (host heap, not pinned) diff --git a/src/stream/config.c b/src/stream/config.c index 8ab9125e..3630e93b 100644 --- a/src/stream/config.c +++ b/src/stream/config.c @@ -139,6 +139,42 @@ compute_level_layout(struct tile_stream_layout* layout, return 1; } +size_t +chunk_scatter_lut_bytes(const struct tile_stream_layout* layout) +{ + const uint64_t columns = layout->input_shape[layout->lifted_rank / 2 - 1]; + const uint64_t rows = layout->epoch_elements / columns; + const size_t limit = SIZE_MAX / sizeof(uint64_t); + if (columns > limit || rows > limit - columns) + return 0; + return (size_t)(columns + rows) * sizeof(uint64_t); +} + +void +chunk_scatter_lut_build(const struct tile_stream_layout* layout, uint64_t* lut) +{ + const int last = layout->lifted_rank / 2 - 1; + const uint64_t columns = layout->input_shape[last]; + for (uint64_t x = 0; x < columns; ++x) + lut[x] = chunk_coordinate_offset(x, + layout->lifted_shape[2 * last + 1], + layout->lifted_strides[2 * last], + layout->lifted_strides[2 * last + 1]); + const uint64_t rows = layout->epoch_elements / columns; + for (uint64_t row = 0; row < rows; ++row) { + uint64_t rest = row, offset = 0; + for (int d = last - 1; d >= 0; --d) { + const uint64_t coord = rest % layout->input_shape[d]; + rest /= layout->input_shape[d]; + offset += chunk_coordinate_offset(coord, + layout->lifted_shape[2 * d + 1], + layout->lifted_strides[2 * d], + layout->lifted_strides[2 * d + 1]); + } + lut[columns + row] = offset; + } +} + // Validate a tile_stream_configuration. // On success, stores the resolved dim partition in *di (slices point into // dims). Returns 0 on success, non-zero on invalid config. @@ -282,6 +318,9 @@ compute_stream_layouts(const struct tile_stream_configuration* config, codec_alignment, storage_order) == 0); } + if (!out->levels.enable_multiscale && + layout_has_partial_chunks(&out->layouts[0])) + CHECK(Fail, chunk_scatter_lut_bytes(&out->layouts[0]) > 0); // --- Level geometry (single loop) --- out->levels.total_chunks = 0; diff --git a/src/stream/layouts.h b/src/stream/layouts.h index 516bd737..eea34129 100644 --- a/src/stream/layouts.h +++ b/src/stream/layouts.h @@ -34,6 +34,23 @@ layout_has_partial_chunks(const struct tile_stream_layout* layout) layout->chunks_per_epoch * layout->chunk_elements; } +#ifdef __cplusplus +extern "C" +{ +#endif + + // A factored logical-to-chunk map: input_shape[last] column offsets followed + // by epoch_elements / input_shape[last] row offsets. All offsets are uint64_t + // elements. This is the same row-plus-column lookup used by LOD chunk + // scatter. Returns 0 bytes if the table size overflows size_t. + size_t chunk_scatter_lut_bytes(const struct tile_stream_layout* layout); + void chunk_scatter_lut_build(const struct tile_stream_layout* layout, + uint64_t* lut); + +#ifdef __cplusplus +} +#endif + // Per-level pre-computed layout information (CPU only, no GPU pointers). struct level_layout_info { diff --git a/src/util/index.ops.h b/src/util/index.ops.h index 026d3bd5..b2b93463 100644 --- a/src/util/index.ops.h +++ b/src/util/index.ops.h @@ -58,6 +58,19 @@ extern "C" int64_t chunk_stride, int64_t* lifted_strides); + // One logical coordinate's contribution to a chunk-pool element offset. +#ifdef __CUDACC__ + __host__ __device__ +#endif + static inline uint64_t chunk_coordinate_offset(uint64_t coord, + uint64_t chunk_size, + int64_t chunk_stride, + int64_t element_stride) + { + return (coord / chunk_size) * (uint64_t)chunk_stride + + (coord % chunk_size) * (uint64_t)element_stride; + } + // Maximum element in shape array. uint64_t max_shape(int ndim, const uint64_t* shape); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e073339d..f8c0bbfe 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -129,10 +129,13 @@ add_gpu_test_exe(test_transpose test_transpose.c LINKS CUDA::cuda_driver CUDA::c # CPU transpose test (no CUDA) add_test_exe(test_transpose_cpu test_transpose_cpu.c transpose_cpu threadpool index_ops_util chucky_log) -add_test_exe(test_packed_input_cpu test_packed_input.c stream_cpu test_shard_sink test_shard_verify) -add_gpu_test_exe(test_packed_input_gpu test_packed_input.c LINKS stream test_shard_sink test_shard_verify CUDA::cuda_driver CUDA::cudart_static) +add_test_exe(test_packed_input_cpu test_packed_input.c stream_cpu multiarray_cpu test_shard_sink test_shard_verify) +add_test_exe(test_scatter_lut_cpu test_scatter_lut.c lod_cpu index_ops_util) +add_gpu_test_exe(test_scatter_lut_gpu test_scatter_lut.c LINKS lod index_ops_util CUDA::cuda_driver CUDA::cudart_static) +add_gpu_test_exe(test_packed_input_gpu test_packed_input.c LINKS stream multiarray_gpu test_shard_sink test_shard_verify CUDA::cuda_driver CUDA::cudart_static) if(CHUCKY_ENABLE_GPU) target_compile_definitions(test_packed_input_gpu PRIVATE TEST_GPU) + target_compile_definitions(test_scatter_lut_gpu PRIVATE TEST_GPU) endif() # CPU compress test (no CUDA) diff --git a/tests/test_packed_input.c b/tests/test_packed_input.c index b54366a1..3d6062f0 100644 --- a/tests/test_packed_input.c +++ b/tests/test_packed_input.c @@ -5,6 +5,7 @@ #include "util/prelude.h" #ifdef TEST_GPU +#include "multiarray.gpu.h" #include "stream.gpu.h" #include "test_runner.h" #define stream_create tile_stream_gpu_create @@ -12,13 +13,26 @@ #define stream_cursor tile_stream_gpu_cursor #define stream_destroy tile_stream_gpu_destroy #define stream_type tile_stream_gpu +#define stream_memory_info tile_stream_memory_info +#define stream_memory_estimate tile_stream_gpu_memory_estimate +#define multi_create multiarray_tile_stream_gpu_create +#define multi_writer multiarray_tile_stream_gpu_writer +#define multi_destroy multiarray_tile_stream_gpu_destroy +#define multi_type multiarray_tile_stream_gpu #else +#include "multiarray.cpu.h" #include "stream.cpu.h" #define stream_create tile_stream_cpu_create #define stream_writer tile_stream_cpu_writer #define stream_cursor tile_stream_cpu_cursor #define stream_destroy tile_stream_cpu_destroy #define stream_type tile_stream_cpu +#define stream_memory_info tile_stream_cpu_memory_info +#define stream_memory_estimate tile_stream_cpu_memory_estimate +#define multi_create multiarray_tile_stream_cpu_create +#define multi_writer multiarray_tile_stream_cpu_writer +#define multi_destroy multiarray_tile_stream_cpu_destroy +#define multi_type multiarray_tile_stream_cpu #endif #include @@ -97,6 +111,14 @@ run_case(uint64_t frames, .max_threads = 4, .reduce_method = lod_reduce_max, }; + if (!multiscale) { + struct stream_memory_info memory; + const uint64_t rows = depth * height * (depth > 1 ? channels : 1); + const size_t table_bytes = + (height % 8 || width % 8) ? (rows + width) * sizeof(uint64_t) : 0; + CHECK(Cleanup, !stream_memory_estimate(&config, 0, &memory)); + CHECK(Cleanup, memory.nlod == 1 && memory.lod_bytes == table_bytes); + } stream = stream_create(&config, &sink.base); CHECK(Cleanup, stream); struct writer* writer = stream_writer(stream); @@ -203,6 +225,103 @@ run_case(uint64_t frames, return error; } +static int +test_multiarray_packed(void) +{ + int error = 1; + struct multi_type* stream = NULL; + struct test_shard_sink sink[3]; + struct shard_sink* sinks[3]; + struct dimension dims[3][3]; + struct tile_stream_configuration configs[3]; + const uint64_t heights[] = { 3, 5, 4 }, widths[] = { 5, 7, 8 }; + // Two different packed mappings and one aligned layout share the pools. + for (int a = 0; a < 3; ++a) { + test_sink_init(&sink[a], 4, 4096); + sinks[a] = &sink[a].base; + dims_create(dims[a], "tyx", (uint64_t[]){ 0, heights[a], widths[a] }); + dims_set_chunk_sizes(dims[a], 3, (uint64_t[]){ 1, 4, 4 }); + dims_set_shard_counts(dims[a], 3, (uint64_t[]){ 0, 1, 1 }); + dims[a][0].chunks_per_shard = 2; + if (a == 1) { + dims[a][1].storage_position = 2; + dims[a][2].storage_position = 1; + } + configs[a] = (struct tile_stream_configuration){ + .buffer_capacity_bytes = 128, + .dtype = dtype_u16, + .rank = 3, + .dimensions = dims[a], + .codec = { .id = CODEC_NONE }, + .epochs_per_batch = 2, + .max_threads = 2, + }; + } + stream = multi_create(3, configs, sinks, 0); + CHECK(Cleanup, stream); + struct multiarray_writer* writer = multi_writer(stream); + for (uint64_t t = 0; t < 5; ++t) + for (int a = 0; a < 3; ++a) { + uint16_t input[35]; + const size_t n = heights[a] * widths[a]; + for (uint64_t y = 0; y < heights[a]; ++y) + for (uint64_t x = 0; x < widths[a]; ++x) + input[y * widths[a] + x] = pixel(t + 10 * a, y, x, widths[a]); + for (size_t begin = 0; begin < n;) { + const size_t end = begin == 0 ? n / 2 : n; + const struct multiarray_writer_result r = writer->update( + writer, a, (struct slice){ input + begin, input + end }); + CHECK(Cleanup, !r.error && r.rest.beg == r.rest.end); + begin = end; + } + } + CHECK(Cleanup, !writer->flush(writer).error && !writer->close(writer).error); + for (int a = 0; a < 3; ++a) { + CHECK(Cleanup, sink[a].last_append_size0 == 5); + const uint64_t cy = ceildiv(heights[a], 4), cx = ceildiv(widths[a], 4); + const size_t slots = 2 * cy * cx; + for (uint64_t sh = 0; sh < 3; ++sh) { + uint64_t offsets[8], sizes[8]; + const struct test_shard_writer* w = &sink[a].writers[0][sh]; + CHECK(Cleanup, w->finalized); + CHECK(Cleanup, !shard_index_check_crc(w->buf, w->size, slots)); + CHECK(Cleanup, + !shard_index_parse(w->buf, w->size, slots, offsets, sizes)); + for (uint64_t t = 0; t < 2; ++t) + for (uint64_t iy = 0; iy < cy; ++iy) + for (uint64_t ix = 0; ix < cx; ++ix) { + const uint64_t slot = + t * cy * cx + (a == 1 ? ix * cy + iy : iy * cx + ix); + if (sh * 2 + t >= 5) { + CHECK(Cleanup, sizes[slot] == UINT64_MAX); + continue; + } + CHECK(Cleanup, + sizes[slot] == 32 && offsets[slot] <= w->size && + sizes[slot] <= w->size - offsets[slot]); + for (uint64_t y = 0; y < 4; ++y) + for (uint64_t x = 0; x < 4; ++x) { + const uint64_t gy = iy * 4 + y, gx = ix * 4 + x; + uint16_t expected = + gy < heights[a] && gx < widths[a] + ? pixel(sh * 2 + t + 10 * a, gy, gx, widths[a]) + : 0; + uint16_t actual; + const uint64_t pos = a == 1 ? x * 4 + y : y * 4 + x; + memcpy(&actual, w->buf + offsets[slot] + pos * 2, 2); + CHECK(Cleanup, actual == expected); + } + } + } + } + error = 0; +Cleanup: + multi_destroy(stream); + for (int a = 0; a < 3; ++a) + test_sink_free(&sink[a]); + return error; +} + static int test_packed(void) { @@ -221,6 +340,7 @@ test_packed(void) CHECK(Fail, !run_case(3, 9, 13, 1, 131, 0, 1, lod, 3)); CHECK(Fail, !run_case(3, 9, 13, 2, 131, 1, 1, lod, 3)); } + CHECK(Fail, !test_multiarray_packed()); return 0; Fail: return 1; diff --git a/tests/test_scatter_lut.c b/tests/test_scatter_lut.c new file mode 100644 index 00000000..d4c7f583 --- /dev/null +++ b/tests/test_scatter_lut.c @@ -0,0 +1,225 @@ +#include "index.ops.util.h" +#include "util/prelude.h" + +#ifdef TEST_GPU +#include "gpu/lod.h" +#include "test_runner.h" +#else +#include "cpu/lod.h" +#endif + +#include +#include + +// Independent row-major chunk and within-chunk coordinates in storage order. +static uint64_t +expected_offset(int rank, + const uint64_t* shape, + const uint64_t* chunks, + const uint8_t* order, + uint64_t chunk_stride, + uint64_t index) +{ + uint64_t coords[HALF_MAX_RANK]; + for (int d = rank - 1; d >= 0; --d) { + coords[d] = index % shape[d]; + index /= shape[d]; + } + uint64_t chunk = 0, within = 0; + for (int j = 0; j < rank; ++j) { + const int d = order ? order[j] : j; + chunk = chunk * ceildiv(shape[d], chunks[d]) + coords[d] / chunks[d]; + within = within * chunks[d] + coords[d] % chunks[d]; + } + return chunk * chunk_stride + within; +} + +static int +run_case(uint8_t rank, + const uint64_t* shape, + const uint64_t* chunks, + const uint8_t* order, + uint8_t bpe) +{ + int error = 1; + struct tile_stream_layout layout; + uint64_t* lut = NULL; + unsigned char *src = NULL, *dst = NULL, *expected = NULL; +#ifdef TEST_GPU + CUdeviceptr d_src = 0, d_dst = 0, d_lut = 0; + CUstream stream = 0; +#endif + CHECK(Fail, + !test_level_layout( + &layout, rank, 1, shape, chunks, order, bpe, TEST_CHUNK_ALIGNMENT)); + const uint64_t epoch = layout.epoch_elements; + const uint64_t width = shape[rank - 1]; + // Start mid-row just before an epoch boundary, then cross several epochs. + const uint64_t start = epoch > 5 ? epoch - 5 : 1; + const uint64_t count = 2 * epoch + 17; + const uint64_t stride = layout.chunks_per_epoch * layout.chunk_stride + 11; + const uint64_t regions = ceildiv(start + count, epoch); + const size_t dst_bytes = regions * stride * bpe; + const size_t lut_bytes = chunk_scatter_lut_bytes(&layout); + lut = malloc(lut_bytes); + src = malloc(count * bpe); + dst = calloc(1, dst_bytes); + expected = calloc(1, dst_bytes); + CHECK(Fail, lut && src && dst && expected); + chunk_scatter_lut_build(&layout, lut); + + for (uint64_t i = 0; i < epoch; ++i) + CHECK( + Fail, + lut[i % width] + lut[width + i / width] == + expected_offset(rank, shape, chunks, order, layout.chunk_stride, i)); + + for (uint64_t i = 0; i < count * bpe; ++i) + src[i] = (unsigned char)(1 + i % 251); + for (uint64_t i = 0; i < count; ++i) { + const uint64_t index = start + i; + const uint64_t offset = + (index / epoch) * stride + + expected_offset( + rank, shape, chunks, order, layout.chunk_stride, index % epoch); + memcpy(expected + offset * bpe, src + i * bpe, bpe); + } + +#ifdef TEST_GPU + CU(Fail, cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING)); + CU(Fail, cuMemAlloc(&d_src, count * bpe)); + CU(Fail, cuMemAlloc(&d_dst, dst_bytes)); + CU(Fail, cuMemAlloc(&d_lut, lut_bytes)); + CU(Fail, cuMemcpyHtoDAsync(d_src, src, count * bpe, stream)); + CU(Fail, cuMemcpyHtoDAsync(d_lut, lut, lut_bytes, stream)); + CU(Fail, cuMemsetD8Async(d_dst, 0, dst_bytes, stream)); + CHECK(Fail, + !scatter_lut_gpu(d_dst, + d_src, + count, + bpe, + start, + width, + epoch, + stride * bpe, + d_lut, + d_lut + width * sizeof(uint64_t), + stream)); + CU(Fail, cuStreamSynchronize(stream)); + CU(Fail, cuMemcpyDtoH(dst, d_dst, dst_bytes)); +#else + for (uint64_t i = 0; i < count;) { + const uint64_t index = start + i; + uint64_t n = epoch - index % epoch; + if (n > count - i) + n = count - i; + CHECK(Fail, + !scatter_lut_cpu(dst + (index / epoch) * stride * bpe, + src + i * bpe, + n, + bpe, + index % epoch, + width, + lut, + lut + width, + NULL)); + i += n; + } +#endif + CHECK(Fail, !memcmp(dst, expected, dst_bytes)); +#ifndef TEST_GPU + // Reuse a partially written epoch, clearing only the mapped logical tail. + CHECK(Fail, + !scatter_lut_cpu( + dst, NULL, epoch - start, bpe, start, width, lut, lut + width, NULL)); + for (uint64_t i = start; i < epoch; ++i) { + const uint64_t offset = + expected_offset(rank, shape, chunks, order, layout.chunk_stride, i); + memset(expected + offset * bpe, 0, bpe); + } + CHECK(Fail, !memcmp(dst, expected, dst_bytes)); +#endif + error = 0; +Fail: +#ifdef TEST_GPU + if (stream) + cuStreamSynchronize(stream); + cuMemFree(d_src); + cuMemFree(d_dst); + cuMemFree(d_lut); + cuStreamDestroy(stream); +#endif + free(lut); + free(src); + free(dst); + free(expected); + return error; +} + +static int +test_scatter_lut(void) +{ + for (uint8_t bpe = 1; bpe <= 8; bpe *= 2) { + CHECK(Fail, + !run_case( + 3, (uint64_t[]){ 2, 9, 13 }, (uint64_t[]){ 2, 8, 8 }, NULL, bpe)); + CHECK(Fail, + !run_case(3, + (uint64_t[]){ 2, 9, 13 }, + (uint64_t[]){ 2, 8, 8 }, + (uint8_t[]){ 0, 2, 1 }, + bpe)); + CHECK(Fail, + !run_case( + 3, (uint64_t[]){ 2, 9, 1 }, (uint64_t[]){ 2, 8, 4 }, NULL, bpe)); + CHECK(Fail, + !run_case( + 3, (uint64_t[]){ 1, 1, 13 }, (uint64_t[]){ 1, 1, 8 }, NULL, bpe)); + CHECK(Fail, !run_case(1, (uint64_t[]){ 7 }, (uint64_t[]){ 7 }, NULL, bpe)); + CHECK(Fail, + !run_case(10, + (uint64_t[]){ 2, 1, 1, 1, 1, 1, 1, 1, 3, 5 }, + (uint64_t[]){ 2, 1, 1, 1, 1, 1, 1, 1, 2, 4 }, + (uint8_t[]){ 0, 1, 2, 3, 4, 5, 6, 7, 9, 8 }, + bpe)); + } + return 0; +Fail: + return 1; +} + +static int +test_wide_offsets(void) +{ + // Wide chunk/element strides with only eight LUT entries: no huge pool. + struct tile_stream_layout layout = { + .lifted_rank = 6, + .lifted_shape = { 1, 1, 2, 2, 2, 4 }, + .lifted_strides = { 0, 0, 1ll << 34, 1ll << 33, 1ll << 35, 3 }, + .input_shape = { 1, 3, 5 }, + .epoch_elements = 15, + }; + uint64_t lut[8]; + const uint64_t expected[] = { 0, 3, 6, 9, + 1ull << 35, 0, 1ull << 33, 1ull << 34 }; + CHECK(Fail, chunk_scatter_lut_bytes(&layout) == sizeof(lut)); + chunk_scatter_lut_build(&layout, lut); + CHECK(Fail, !memcmp(lut, expected, sizeof(lut))); + layout.lifted_rank = 2; + layout.input_shape[0] = layout.epoch_elements = SIZE_MAX / sizeof(uint64_t); + CHECK(Fail, chunk_scatter_lut_bytes(&layout) == 0); + return 0; +Fail: + return 1; +} + +#ifdef TEST_GPU +RUN_GPU_TESTS({ "scatter_lut", test_scatter_lut }, + { "wide_offsets", test_wide_offsets }) +#else +int +main(void) +{ + return test_scatter_lut() | test_wide_offsets(); +} +#endif From 24432e163fd2da15694a4045a2f974bf48c268a1 Mon Sep 17 00:00:00 2001 From: Nathan Clack Date: Sun, 4 Oct 2026 11:55:25 -0700 Subject: [PATCH 4/4] Cache contiguous runs for packed CPU scatter Compile column spans and joined rows from the lookup tables once per array. Use those boundaries for bulk copies and tail clearing while retaining typed scatter for short spans. Reuse a single column boundary when padded rows cannot join, and include a conservative cache allowance in memory estimates. Cover partial appends, worker splits, row padding, ragged multidimensional spans, cache replacement, and multiarray switches with independent readback. --- README.md | 2 + src/cpu/lod.cpp | 134 +++++++++++++++++++++++++++++++++++- src/cpu/lod.h | 22 +++++- src/cpu/stream.body.c | 1 + src/cpu/stream.body.h | 1 + src/cpu/stream.c | 19 +++++ src/cpu/stream.internal.h | 2 + src/multiarray/stream.c | 12 ++++ src/stream.cpu.h | 3 +- tests/test_packed_input.c | 121 +++++++++++++++++++++----------- tests/test_scatter_lut.c | 141 +++++++++++++++++++++++++++++++++----- 11 files changed, 393 insertions(+), 65 deletions(-) diff --git a/README.md b/README.md index 13718376..767402e4 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,8 @@ it alongside L0. Input remains tightly packed. At partial chunk edges, precomputed row and column offsets map logical coordinates into zeroed output chunks through the same lookup scatter used for LOD. Padding requires no work from the caller. +CPU scatter caches contiguous run boundaries so rows and groups of rows can +be copied in bulk. **Supported element types:** u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64 (see `enum dtype` in `src/dtype.h`). diff --git a/src/cpu/lod.cpp b/src/cpu/lod.cpp index 02f16066..e762f585 100644 --- a/src/cpu/lod.cpp +++ b/src/cpu/lod.cpp @@ -7,6 +7,7 @@ extern "C" #include "util/prelude.h" } +#include #include #include #include @@ -387,6 +388,93 @@ reduce_typed(const lod_plan* p, } } +// Small memcpy calls lose to typed stores in the packed-input benchmarks. +// This is a per-run choice, independent of thread-pool dispatch size. +static constexpr size_t min_bulk_copy_bytes = 64; + +struct scatter_run_stats +{ + size_t count; + uint64_t longest; +}; + +static scatter_run_stats +find_scatter_runs(const uint64_t* offsets, + uint64_t count, + uint64_t stride, + uint64_t unit, + uint64_t* ends) +{ + scatter_run_stats stats = {}; + uint64_t begin = 0; + for (uint64_t i = 1; i <= count; ++i) { + if (i != count && offsets[i] >= offsets[i - 1] && + offsets[i] - offsets[i - 1] == stride) + continue; + if (ends) + ends[stats.count] = i * unit; + ++stats.count; + stats.longest = std::max(stats.longest, (i - begin) * unit); + begin = i; + } + return stats; +} + +extern "C" int +scatter_lut_runs_build(struct scatter_lut_runs* runs, + uint8_t bpe, + uint64_t width, + uint64_t row_count, + const uint64_t* columns, + const uint64_t* rows) +{ + if ((bpe != 1 && bpe != 2 && bpe != 4 && bpe != 8) || !width || !row_count || + row_count > UINT64_MAX / width || width > SIZE_MAX / sizeof(uint64_t) || + row_count > SIZE_MAX / sizeof(uint64_t)) + return 1; + + const uint64_t* offsets = columns; + uint64_t count = width, stride = 1, unit = 1; + scatter_run_stats stats = + find_scatter_runs(offsets, count, stride, unit, NULL); + if (stats.count == 1) { + // A contiguous column table permits joining whole rows. Row offsets + // include padding, chunk boundaries, and every remaining dimension. + const auto row_stats = + find_scatter_runs(rows, row_count, width, width, NULL); + if (row_stats.count < row_count) { + offsets = rows; + count = row_count; + stride = unit = width; + stats = row_stats; + } + // If no rows join, reuse one column boundary for every row instead of + // caching identical row lengths across the entire epoch. + } + + scatter_lut_runs next = {}; + if (stats.longest >= min_bulk_copy_bytes / bpe) { + if (stats.count > SIZE_MAX / sizeof(uint64_t)) + return 1; + next.ends = (uint64_t*)malloc(stats.count * sizeof(uint64_t)); + if (!next.ends) + return 1; + next.count = stats.count; + next.period = count * unit; + find_scatter_runs(offsets, count, stride, unit, next.ends); + } + scatter_lut_runs_free(runs); + *runs = next; + return 0; +} + +extern "C" void +scatter_lut_runs_free(struct scatter_lut_runs* runs) +{ + free(runs->ends); + *runs = {}; +} + template static void scatter_lut_typed(const T* values, @@ -396,10 +484,11 @@ scatter_lut_typed(const T* values, uint64_t width, uint64_t count, uint64_t i_offset, + const struct scatter_lut_runs* runs, struct threadpool* pool) { // Decompose once per row, including ranges split by the thread pool. - auto body = [=](size_t beg, size_t end) { + auto scatter = [=](size_t beg, size_t end) { while (beg < end) { const uint64_t index = i_offset + beg; const uint64_t column = index % width; @@ -411,6 +500,36 @@ scatter_lut_typed(const T* values, beg += n; } }; + auto body = [=](size_t beg, size_t end) { + if (!runs || !runs->count) { + scatter(beg, end); + return; + } + uint64_t position = (i_offset + beg) % runs->period; + const uint64_t* run = + std::upper_bound(runs->ends, runs->ends + runs->count, position); + while (beg < end) { + const uint64_t index = i_offset + beg; + const uint64_t n = std::min(*run - position, end - beg); + T* dst = chunks + batch_offsets[index / width] + chunk_lut[index % width]; + if (n >= min_bulk_copy_bytes / sizeof(T)) { + if (values) + memcpy(dst, values + beg, n * sizeof(T)); + else + memset(dst, 0, n * sizeof(T)); + } else { + for (uint64_t i = 0; i < n; ++i) + dst[i] = values ? values[beg + i] : T{}; + } + beg += n; + position += n; + ++run; + if (position == runs->period) { + position = 0; + run = runs->ends; + } + } + }; auto trampoline = [](size_t beg, size_t end, int tid, void* ctx) { (void)tid; (*static_cast(ctx))(beg, end); @@ -430,12 +549,20 @@ scatter_lut_cpu(void* dst, uint64_t width, const uint64_t* columns, const uint64_t* rows, + const struct scatter_lut_runs* runs, struct threadpool* pool) { #define CASE(B, T) \ case B: \ - scatter_lut_typed( \ - (const T*)src, (T*)dst, columns, rows, width, count, i_offset, pool); \ + scatter_lut_typed((const T*)src, \ + (T*)dst, \ + columns, \ + rows, \ + width, \ + count, \ + i_offset, \ + runs, \ + pool); \ return 0 switch (bpe) { CASE(1, uint8_t); @@ -652,6 +779,7 @@ lod_cpu_morton_to_chunks(const lod_plan* p, lod_count, \ p->levels.level[lv].fixed_dims_count* lod_count, \ 0, \ + NULL, \ pool) DISPATCH(dtype, DO); #undef DO diff --git a/src/cpu/lod.h b/src/cpu/lod.h index f6a96f22..53f32518 100644 --- a/src/cpu/lod.h +++ b/src/cpu/lod.h @@ -16,9 +16,28 @@ extern "C" struct threadpool; + // Cached contiguous run ends, in input elements. The runs repeat every + // period elements: one column table, or a whole epoch when rows coalesce. + // An empty cache selects element scatter. Owns ends; initialize to zero. + struct scatter_lut_runs + { + uint64_t* ends; + size_t count; + uint64_t period; + }; + + // Derive runs from the LUT once. Short-only mappings need no allocation. + int scatter_lut_runs_build(struct scatter_lut_runs* runs, + uint8_t bpe, + uint64_t width, + uint64_t row_count, + const uint64_t* columns, + const uint64_t* rows); + void scatter_lut_runs_free(struct scatter_lut_runs* runs); + // Scatter a source range through a row-plus-column lookup. i_offset is // within one epoch; source starts at the supplied range. NULL src clears - // the mapped elements (used for the final partial epoch). + // the mapped elements (used for the final partial epoch). runs may be NULL. int scatter_lut_cpu(void* dst, const void* src, uint64_t count, @@ -27,6 +46,7 @@ extern "C" uint64_t width, const uint64_t* columns, const uint64_t* rows, + const struct scatter_lut_runs* runs, struct threadpool* pool); // Reduce across LOD levels in-place. diff --git a/src/cpu/stream.body.c b/src/cpu/stream.body.c index 4ace6243..be00e517 100644 --- a/src/cpu/stream.body.c +++ b/src/cpu/stream.body.c @@ -110,6 +110,7 @@ scatter_input(struct cpu_stream_view* v, width, v->input_chunk_lut, v->input_chunk_lut + width, + v->input_chunk_runs, v->pool); } return transpose_cpu(dst, src, bytes, bpe, i_offset, v->layout, v->pool); diff --git a/src/cpu/stream.body.h b/src/cpu/stream.body.h index 1c7be041..13900c07 100644 --- a/src/cpu/stream.body.h +++ b/src/cpu/stream.body.h @@ -32,6 +32,7 @@ struct cpu_stream_view struct aggregate_layout* agg_layout; // [LOD_MAX_LEVELS] array struct reduce_csr* csrs; // [nlod-1] CSR LUTs const uint64_t* input_chunk_lut; // direct row/column map, or NULL + const struct scatter_lut_runs* input_chunk_runs; void* append_accum; uint32_t* append_counts; // [LOD_MAX_LEVELS] diff --git a/src/cpu/stream.c b/src/cpu/stream.c index 94857670..f120b777 100644 --- a/src/cpu/stream.c +++ b/src/cpu/stream.c @@ -70,6 +70,14 @@ tile_stream_cpu_create(const struct tile_stream_configuration* config, s->input_chunk_lut = (uint64_t*)malloc(bytes); CHECK(Fail, s->input_chunk_lut); chunk_scatter_lut_build(&s->layout, s->input_chunk_lut); + const uint64_t width = s->layout.input_shape[config->rank - 1]; + CHECK(Fail, + !scatter_lut_runs_build(&s->input_chunk_runs, + (uint8_t)dtype_bpe(config->dtype), + width, + s->layout.epoch_elements / width, + s->input_chunk_lut, + s->input_chunk_lut + width)); } const uint32_t K = s->cl.epochs_per_batch; @@ -359,6 +367,7 @@ tile_stream_cpu_release_resources(struct tile_stream_cpu* s) free(s->batch_active_masks); free(s->pool_epochs_scratch); free(s->input_chunk_lut); + scatter_lut_runs_free(&s->input_chunk_runs); free(s->scatter_lut); free(s->scatter_fixed_dims_offsets); free(s->linear); @@ -535,6 +544,15 @@ compute_memory_info(const struct computed_stream_layouts* cl, layout_has_partial_chunks(&cl->layouts[0])) { lod = chunk_scatter_lut_bytes(&cl->layouts[0]); CHECK(Error, lod); + const struct tile_stream_layout* layout = &cl->layouts[0]; + const uint64_t width = layout->input_shape[cl->rank - 1]; + const uint64_t rows = layout->epoch_elements / width; + // The cache has at most one end per column or per row. Reserve that + // upper bound without allocating/scanning the LUT during estimation. + const size_t cache = + (size_t)(width > rows ? width : rows) * sizeof(uint64_t); + CHECK(Error, cache <= SIZE_MAX - lod); + lod += cache; } if (cl->levels.enable_multiscale) { lod += cl->layouts[0].epoch_elements * bytes_per_element; // linear @@ -808,6 +826,7 @@ make_view(struct tile_stream_cpu* s) .pool_epochs_scratch = s->pool_epochs_scratch, .pool_fully_covered = s->pool_fully_covered, .input_chunk_lut = s->input_chunk_lut, + .input_chunk_runs = &s->input_chunk_runs, .shard = s->shard, .agg_layout = s->agg_layout, .csrs = s->csrs, diff --git a/src/cpu/stream.internal.h b/src/cpu/stream.internal.h index 6a868b6f..1c340a2f 100644 --- a/src/cpu/stream.internal.h +++ b/src/cpu/stream.internal.h @@ -1,5 +1,6 @@ #pragma once +#include "cpu/lod.h" #include "cpu/pipeline.h" #include "cpu/stream.body.h" #include "lod/reduce_csr.h" @@ -21,6 +22,7 @@ struct tile_stream_cpu struct level_geometry levels; uint64_t* input_chunk_lut; // direct row/column map (single-level edges) + struct scatter_lut_runs input_chunk_runs; // Chunk pool: total_chunks * chunk_stride * bpe bytes. void* chunk_pool; diff --git a/src/multiarray/stream.c b/src/multiarray/stream.c index 9c84ea80..8f783f70 100644 --- a/src/multiarray/stream.c +++ b/src/multiarray/stream.c @@ -10,6 +10,7 @@ #include "cpu/compress.h" #include "cpu/compress_blosc.h" +#include "cpu/lod.h" #include "cpu/transpose.h" #include "util/metric.h" #include "util/prelude.h" @@ -40,6 +41,7 @@ struct array_descriptor struct computed_stream_layouts cl; struct tile_stream_layout layout; uint64_t* input_chunk_lut; + struct scatter_lut_runs input_chunk_runs; struct level_geometry levels; struct aggregate_layout agg_layout[LOD_MAX_LEVELS]; struct shard_state shard[LOD_MAX_LEVELS]; @@ -184,6 +186,14 @@ init_array_descriptor(struct array_descriptor* desc, if (!bytes || !(desc->input_chunk_lut = (uint64_t*)malloc(bytes))) return 1; chunk_scatter_lut_build(&desc->layout, desc->input_chunk_lut); + const uint64_t width = desc->layout.input_shape[config->rank - 1]; + if (scatter_lut_runs_build(&desc->input_chunk_runs, + (uint8_t)dtype_bpe(config->dtype), + width, + desc->layout.epoch_elements / width, + desc->input_chunk_lut, + desc->input_chunk_lut + width)) + return 1; } desc->pool_fully_covered = (desc->layout.chunk_stride == desc->layout.chunk_elements); @@ -526,6 +536,7 @@ multiarray_tile_stream_cpu_release_resources( free(desc->batch_active_masks); free(desc->pool_epochs_scratch); free(desc->input_chunk_lut); + scatter_lut_runs_free(&desc->input_chunk_runs); if (desc->csrs) { int ncsr = desc->cl.plan.levels.nlod - 1; for (int l = 0; l < ncsr; ++l) @@ -680,6 +691,7 @@ make_multiarray_view(struct multiarray_tile_stream_cpu* ms, .pool_epochs_scratch = desc->pool_epochs_scratch, .pool_fully_covered = desc->pool_fully_covered, .input_chunk_lut = desc->input_chunk_lut, + .input_chunk_runs = &desc->input_chunk_runs, .shard = desc->shard, .agg_layout = desc->agg_layout, .csrs = desc->csrs, diff --git a/src/stream.cpu.h b/src/stream.cpu.h index f464f8a8..49c90012 100644 --- a/src/stream.cpu.h +++ b/src/stream.cpu.h @@ -50,7 +50,7 @@ struct tile_stream_cpu_memory_info size_t comp_sizes_bytes; // K * total_chunks * sizeof(size_t) size_t aggregate_bytes; // 2x per-batch scratch slots + batch masks size_t host_output_pool_bytes; - size_t lod_bytes; // scatter LUTs and multiscale buffers/accumulators + size_t lod_bytes; // scatter LUTs/run-cache bound and multiscale buffers size_t shard_bytes; // active_shard arrays + index buffers uint64_t chunks_per_epoch; @@ -62,6 +62,7 @@ struct tile_stream_cpu_memory_info }; // shard_alignment: 0 = no alignment constraint. +// Packed input includes a conservative allowance for cached scatter run ends. int tile_stream_cpu_memory_estimate(const struct tile_stream_configuration* config, size_t shard_alignment, diff --git a/tests/test_packed_input.c b/tests/test_packed_input.c index 3d6062f0..98362084 100644 --- a/tests/test_packed_input.c +++ b/tests/test_packed_input.c @@ -45,15 +45,17 @@ pixel(uint64_t t, uint64_t y, uint64_t x, uint64_t width) } static int -run_case(uint64_t frames, - uint64_t height, - uint64_t width, - uint64_t depth, - size_t append_elements, - int bounded, - int transpose, - int multiscale, - uint64_t channels) +run_case_chunks(uint64_t frames, + uint64_t height, + uint64_t width, + uint64_t depth, + size_t append_elements, + int bounded, + int transpose, + int multiscale, + uint64_t channels, + uint64_t chunk_y, + uint64_t chunk_x) { int error = 1; struct stream_type* stream = NULL; @@ -85,13 +87,13 @@ run_case(uint64_t frames, if (rank == 3) { dims_create( dims, "tyx", (uint64_t[]){ bounded ? frames : 0, height, width }); - dims_set_chunk_sizes(dims, 3, (uint64_t[]){ depth, 8, 8 }); + dims_set_chunk_sizes(dims, 3, (uint64_t[]){ depth, chunk_y, chunk_x }); dims_set_shard_counts(dims, 3, (uint64_t[]){ 0, 1, 1 }); } else { dims_create(dims, "tcyx", (uint64_t[]){ bounded ? frames : 0, channels, height, width }); - dims_set_chunk_sizes(dims, 4, (uint64_t[]){ depth, 1, 8, 8 }); + dims_set_chunk_sizes(dims, 4, (uint64_t[]){ depth, 1, chunk_y, chunk_x }); dims_set_shard_counts(dims, 4, (uint64_t[]){ 0, 1, 1, 1 }); } dims[0].chunks_per_shard = 4; @@ -114,10 +116,16 @@ run_case(uint64_t frames, if (!multiscale) { struct stream_memory_info memory; const uint64_t rows = depth * height * (depth > 1 ? channels : 1); - const size_t table_bytes = - (height % 8 || width % 8) ? (rows + width) * sizeof(uint64_t) : 0; + const size_t table_bytes = (height % chunk_y || width % chunk_x) + ? (rows + width) * sizeof(uint64_t) + : 0; + size_t expected_lod_bytes = table_bytes; +#ifndef TEST_GPU + if (table_bytes) + expected_lod_bytes += (rows > width ? rows : width) * sizeof(uint64_t); +#endif CHECK(Cleanup, !stream_memory_estimate(&config, 0, &memory)); - CHECK(Cleanup, memory.nlod == 1 && memory.lod_bytes == table_bytes); + CHECK(Cleanup, memory.nlod == 1 && memory.lod_bytes == expected_lod_bytes); } stream = stream_create(&config, &sink.base); CHECK(Cleanup, stream); @@ -146,7 +154,7 @@ run_case(uint64_t frames, for (int lv = 0; lv < nlod; ++lv) { const uint64_t scale = 1u << lv; const uint64_t ny = ceildiv(height, scale), nx = ceildiv(width, scale); - const uint64_t cy = ceildiv(ny, 8), cx = ceildiv(nx, 8); + const uint64_t cy = ceildiv(ny, chunk_y), cx = ceildiv(nx, chunk_x); const uint64_t ct = bounded && ceildiv(frames, depth) < 4 ? ceildiv(frames, depth) : 4; const size_t slots = ct * channels * cy * cx; @@ -172,7 +180,7 @@ run_case(uint64_t frames, bad = sizes[slot] != UINT64_MAX; continue; } - if (sizes[slot] != depth * 8 * 8 * sizeof(uint16_t) || + if (sizes[slot] != depth * chunk_y * chunk_x * sizeof(uint16_t) || offsets[slot] > w->size || sizes[slot] > w->size - offsets[slot]) { bad = 1; break; @@ -180,9 +188,10 @@ run_case(uint64_t frames, const uint64_t yc = transpose ? slot % cy : (slot / cx) % cy; const uint64_t xc = transpose ? (slot / cy) % cx : slot % cx; for (uint64_t t = 0; t < depth; ++t) - for (uint64_t y = 0; y < 8; ++y) - for (uint64_t x = 0; x < 8; ++x) { - uint64_t gt = tc * depth + t, gy = yc * 8 + y, gx = xc * 8 + x; + for (uint64_t y = 0; y < chunk_y; ++y) + for (uint64_t x = 0; x < chunk_x; ++x) { + uint64_t gt = tc * depth + t, gy = yc * chunk_y + y, + gx = xc * chunk_x + x; uint16_t expected = 0, actual; if (gt < frames && gy < ny && gx < nx) { // Independent max-reduction reference, including odd edges. @@ -198,7 +207,9 @@ run_case(uint64_t frames, expected = value; } } - const uint64_t pos = t * 64 + (transpose ? x * 8 + y : y * 8 + x); + const uint64_t pos = + t * chunk_y * chunk_x + + (transpose ? x * chunk_y + y : y * chunk_x + x); memcpy(&actual, w->buf + offsets[slot] + pos * 2, 2); if (actual != expected && !bad) { log_error("level %d pixel (%llu,%llu,%llu): %u != %u", @@ -225,25 +236,49 @@ run_case(uint64_t frames, return error; } +static int +run_case(uint64_t frames, + uint64_t height, + uint64_t width, + uint64_t depth, + size_t append_elements, + int bounded, + int transpose, + int multiscale, + uint64_t channels) +{ + return run_case_chunks(frames, + height, + width, + depth, + append_elements, + bounded, + transpose, + multiscale, + channels, + 8, + 8); +} + static int test_multiarray_packed(void) { int error = 1; struct multi_type* stream = NULL; - struct test_shard_sink sink[3]; - struct shard_sink* sinks[3]; - struct dimension dims[3][3]; - struct tile_stream_configuration configs[3]; - const uint64_t heights[] = { 3, 5, 4 }, widths[] = { 5, 7, 8 }; - // Two different packed mappings and one aligned layout share the pools. - for (int a = 0; a < 3; ++a) { - test_sink_init(&sink[a], 4, 4096); + struct test_shard_sink sink[4]; + struct shard_sink* sinks[4]; + struct dimension dims[4][3]; + struct tile_stream_configuration configs[4]; + const uint64_t heights[] = { 3, 5, 5, 4 }, widths[] = { 129, 128, 7, 128 }; + // Column runs, joined rows, permuted scatter, and aligned copies share pools. + for (int a = 0; a < 4; ++a) { + test_sink_init(&sink[a], 4, 16384); sinks[a] = &sink[a].base; dims_create(dims[a], "tyx", (uint64_t[]){ 0, heights[a], widths[a] }); - dims_set_chunk_sizes(dims[a], 3, (uint64_t[]){ 1, 4, 4 }); + dims_set_chunk_sizes(dims[a], 3, (uint64_t[]){ 1, 4, 128 }); dims_set_shard_counts(dims[a], 3, (uint64_t[]){ 0, 1, 1 }); dims[a][0].chunks_per_shard = 2; - if (a == 1) { + if (a == 2) { dims[a][1].storage_position = 2; dims[a][2].storage_position = 1; } @@ -257,12 +292,12 @@ test_multiarray_packed(void) .max_threads = 2, }; } - stream = multi_create(3, configs, sinks, 0); + stream = multi_create(4, configs, sinks, 0); CHECK(Cleanup, stream); struct multiarray_writer* writer = multi_writer(stream); for (uint64_t t = 0; t < 5; ++t) - for (int a = 0; a < 3; ++a) { - uint16_t input[35]; + for (int a = 0; a < 4; ++a) { + uint16_t input[640]; const size_t n = heights[a] * widths[a]; for (uint64_t y = 0; y < heights[a]; ++y) for (uint64_t x = 0; x < widths[a]; ++x) @@ -276,9 +311,9 @@ test_multiarray_packed(void) } } CHECK(Cleanup, !writer->flush(writer).error && !writer->close(writer).error); - for (int a = 0; a < 3; ++a) { + for (int a = 0; a < 4; ++a) { CHECK(Cleanup, sink[a].last_append_size0 == 5); - const uint64_t cy = ceildiv(heights[a], 4), cx = ceildiv(widths[a], 4); + const uint64_t cy = ceildiv(heights[a], 4), cx = ceildiv(widths[a], 128); const size_t slots = 2 * cy * cx; for (uint64_t sh = 0; sh < 3; ++sh) { uint64_t offsets[8], sizes[8]; @@ -291,23 +326,23 @@ test_multiarray_packed(void) for (uint64_t iy = 0; iy < cy; ++iy) for (uint64_t ix = 0; ix < cx; ++ix) { const uint64_t slot = - t * cy * cx + (a == 1 ? ix * cy + iy : iy * cx + ix); + t * cy * cx + (a == 2 ? ix * cy + iy : iy * cx + ix); if (sh * 2 + t >= 5) { CHECK(Cleanup, sizes[slot] == UINT64_MAX); continue; } CHECK(Cleanup, - sizes[slot] == 32 && offsets[slot] <= w->size && + sizes[slot] == 1024 && offsets[slot] <= w->size && sizes[slot] <= w->size - offsets[slot]); for (uint64_t y = 0; y < 4; ++y) - for (uint64_t x = 0; x < 4; ++x) { - const uint64_t gy = iy * 4 + y, gx = ix * 4 + x; + for (uint64_t x = 0; x < 128; ++x) { + const uint64_t gy = iy * 4 + y, gx = ix * 128 + x; uint16_t expected = gy < heights[a] && gx < widths[a] ? pixel(sh * 2 + t + 10 * a, gy, gx, widths[a]) : 0; uint16_t actual; - const uint64_t pos = a == 1 ? x * 4 + y : y * 4 + x; + const uint64_t pos = a == 2 ? x * 4 + y : y * 128 + x; memcpy(&actual, w->buf + offsets[slot] + pos * 2, 2); CHECK(Cleanup, actual == expected); } @@ -317,7 +352,7 @@ test_multiarray_packed(void) error = 0; Cleanup: multi_destroy(stream); - for (int a = 0; a < 3; ++a) + for (int a = 0; a < 4; ++a) test_sink_free(&sink[a]); return error; } @@ -340,6 +375,10 @@ test_packed(void) CHECK(Fail, !run_case(3, 9, 13, 1, 131, 0, 1, lod, 3)); CHECK(Fail, !run_case(3, 9, 13, 2, 131, 1, 1, lod, 3)); } + // Cached copies must preserve split appends, partial epochs, and row padding. + CHECK(Fail, !run_case_chunks(5, 9, 257, 2, 131, 1, 0, 0, 1, 8, 128)); + CHECK(Fail, !run_case_chunks(5, 127, 128, 2, 70001, 0, 0, 0, 1, 128, 128)); + CHECK(Fail, !run_case_chunks(5, 127, 127, 2, 509, 1, 0, 0, 1, 128, 128)); CHECK(Fail, !test_multiarray_packed()); return 0; Fail: diff --git a/tests/test_scatter_lut.c b/tests/test_scatter_lut.c index d4c7f583..f8494033 100644 --- a/tests/test_scatter_lut.c +++ b/tests/test_scatter_lut.c @@ -6,6 +6,7 @@ #include "test_runner.h" #else #include "cpu/lod.h" +#include "threadpool/threadpool.h" #endif #include @@ -48,6 +49,9 @@ run_case(uint8_t rank, #ifdef TEST_GPU CUdeviceptr d_src = 0, d_dst = 0, d_lut = 0; CUstream stream = 0; +#else + struct scatter_lut_runs runs = { 0 }; + struct threadpool* pool = NULL; #endif CHECK(Fail, !test_level_layout( @@ -108,36 +112,65 @@ run_case(uint8_t rank, CU(Fail, cuStreamSynchronize(stream)); CU(Fail, cuMemcpyDtoH(dst, d_dst, dst_bytes)); #else - for (uint64_t i = 0; i < count;) { - const uint64_t index = start + i; - uint64_t n = epoch - index % epoch; - if (n > count - i) - n = count - i; - CHECK(Fail, - !scatter_lut_cpu(dst + (index / epoch) * stride * bpe, - src + i * bpe, - n, - bpe, - index % epoch, - width, - lut, - lut + width, - NULL)); - i += n; + CHECK(Fail, + !scatter_lut_runs_build( + &runs, bpe, width, epoch / width, lut, lut + width)); + pool = threadpool_new(3); + CHECK(Fail, pool); + // Both paths must match coordinates, including worker splits inside runs. + for (int cached = 0; cached < 2; ++cached) { + memset(dst, 0, dst_bytes); + for (uint64_t i = 0; i < count;) { + const uint64_t index = start + i; + uint64_t n = epoch - index % epoch; + if (n > count - i) + n = count - i; + CHECK(Fail, + !scatter_lut_cpu(dst + (index / epoch) * stride * bpe, + src + i * bpe, + n, + bpe, + index % epoch, + width, + lut, + lut + width, + cached ? &runs : NULL, + pool)); + i += n; + } + CHECK(Fail, !memcmp(dst, expected, dst_bytes)); } #endif CHECK(Fail, !memcmp(dst, expected, dst_bytes)); #ifndef TEST_GPU // Reuse a partially written epoch, clearing only the mapped logical tail. CHECK(Fail, - !scatter_lut_cpu( - dst, NULL, epoch - start, bpe, start, width, lut, lut + width, NULL)); + !scatter_lut_cpu(dst, + NULL, + epoch - start, + bpe, + start, + width, + lut, + lut + width, + &runs, + pool)); for (uint64_t i = start; i < epoch; ++i) { const uint64_t offset = expected_offset(rank, shape, chunks, order, layout.chunk_stride, i); memset(expected + offset * bpe, 0, bpe); } CHECK(Fail, !memcmp(dst, expected, dst_bytes)); + // Clear a long range crossing cached boundaries and thread-pool slices. + CHECK(Fail, + !scatter_lut_cpu( + dst, NULL, epoch - 1, bpe, 1, width, lut, lut + width, &runs, pool)); + for (uint64_t i = 1; i < epoch; ++i) { + const uint64_t offset = + expected_offset(rank, shape, chunks, order, layout.chunk_stride, i); + memset(expected + offset * bpe, 0, bpe); + } + CHECK(Fail, !memcmp(dst, expected, dst_bytes)); #endif error = 0; Fail: @@ -148,6 +181,9 @@ run_case(uint8_t rank, cuMemFree(d_dst); cuMemFree(d_lut); cuStreamDestroy(stream); +#else + scatter_lut_runs_free(&runs); + threadpool_free(pool); #endif free(lut); free(src); @@ -182,6 +218,38 @@ test_scatter_lut(void) (uint64_t[]){ 2, 1, 1, 1, 1, 1, 1, 1, 2, 4 }, (uint8_t[]){ 0, 1, 2, 3, 4, 5, 6, 7, 9, 8 }, bpe)); + // Whole-frame copies, row padding, long x spans, and ragged y/x blocks. + CHECK(Fail, + !run_case(3, + (uint64_t[]){ 1, 127, 128 }, + (uint64_t[]){ 1, 128, 128 }, + NULL, + bpe)); + CHECK(Fail, + !run_case(3, + (uint64_t[]){ 2, 127, 127 }, + (uint64_t[]){ 2, 128, 128 }, + NULL, + bpe)); + CHECK(Fail, + !run_case(3, + (uint64_t[]){ 2, 257, 263 }, + (uint64_t[]){ 2, 128, 128 }, + NULL, + bpe)); + CHECK(Fail, + !run_case(4, + (uint64_t[]){ 1, 4, 257, 128 }, + (uint64_t[]){ 1, 2, 128, 128 }, + NULL, + bpe)); + // Equal first/last offsets do not prove interior rows are contiguous. + CHECK(Fail, + !run_case(4, + (uint64_t[]){ 1, 4, 4, 64 }, + (uint64_t[]){ 1, 4, 4, 64 }, + (uint8_t[]){ 0, 2, 1, 3 }, + bpe)); } return 0; Fail: @@ -213,6 +281,41 @@ test_wide_offsets(void) return 1; } +#ifndef TEST_GPU +static int +test_cached_runs(void) +{ + int error = 1; + struct scatter_lut_runs runs = { 0 }; + uint64_t columns[129], rows[] = { 0, 256, 512 }; + for (uint64_t i = 0; i < 128; ++i) + columns[i] = i; + columns[128] = 1024; + + // Long x spans plus a short edge: the cache must actually be populated. + CHECK(Fail, !scatter_lut_runs_build(&runs, 1, 129, 3, columns, rows)); + CHECK(Fail, runs.count == 2 && runs.period == 129); + CHECK(Fail, runs.ends[0] == 128 && runs.ends[1] == 129); + + // Padded rows reuse one boundary; rebuilding releases the previous cache. + CHECK(Fail, !scatter_lut_runs_build(&runs, 1, 128, 3, columns, rows)); + CHECK(Fail, runs.count == 1 && runs.period == 128 && runs.ends[0] == 128); + + rows[1] = 128; + rows[2] = 256; + CHECK(Fail, !scatter_lut_runs_build(&runs, 1, 128, 3, columns, rows)); + CHECK(Fail, runs.count == 1 && runs.period == 384 && runs.ends[0] == 384); + + // Tiny isolated spans retain element scatter and allocate no boundaries. + CHECK(Fail, !scatter_lut_runs_build(&runs, 1, 8, 3, columns, rows)); + CHECK(Fail, !runs.count && !runs.ends); + error = 0; +Fail: + scatter_lut_runs_free(&runs); + return error; +} +#endif + #ifdef TEST_GPU RUN_GPU_TESTS({ "scatter_lut", test_scatter_lut }, { "wide_offsets", test_wide_offsets }) @@ -220,6 +323,6 @@ RUN_GPU_TESTS({ "scatter_lut", test_scatter_lut }, int main(void) { - return test_scatter_lut() | test_wide_offsets(); + return test_scatter_lut() | test_wide_offsets() | test_cached_runs(); } #endif