From e3b33591ca765cc1065d76015ad628718e5fd37f Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 14 Oct 2026 21:49:37 +0200 Subject: [PATCH] GPU: kernels read host buffers in place on unified-memory systems gpu_direct_host_read() (CUDAWrapper) is true where the device accesses pageable memory through the host's page tables (cudaDevAttrPageableMemoryAccess and ...UsesHostPageTables: GB10, GH200, GB200), asked once per device; JFJOCH_FORCE_DIRECT_HOST_READ forces it on for testing with page-locked inputs. Where it holds: - BSLZ4DecoderGPU: the LZ4/fused kernels and the tail copy read the reader's compressed chunk directly, no pinned staging memcpy and no H2D copy; - ImagePreprocessorGPU: Analyze and the UINT16_WIDENED path hand the host buffer to preprocess_kernel (if aligned to the pixel type), and nothing is page-locked; - RotationScaleMergeGPU::IngestUpload: IngestExtractKernel reads the pinned stage in place, without device staging buffers; the stage is reused only after the kernel has run. Discrete cards take the unchanged path. Every caller synchronises its stream before the host buffer goes away. DirectHostReadGPUTest checks the three sites against the CPU on page-locked inputs, in either mode. Co-Authored-By: Claude Opus 5.5 --- common/CUDAWrapper.cpp | 4 + common/CUDAWrapper.cu | 23 +++ common/CUDAWrapper.h | 8 + .../image_preprocessing/BSLZ4DecoderGPU.cu | 24 ++- .../image_preprocessing/BSLZ4DecoderGPU.h | 3 + .../ImagePreprocessorGPU.cu | 44 ++++-- .../ImagePreprocessorGPU.h | 12 +- .../scale_merge/RotationScaleMergeGPU.cu | 25 ++- tests/CMakeLists.txt | 1 + tests/DirectHostReadGPUTest.cpp | 148 ++++++++++++++++++ 10 files changed, 263 insertions(+), 29 deletions(-) create mode 100644 tests/DirectHostReadGPUTest.cpp diff --git a/common/CUDAWrapper.cpp b/common/CUDAWrapper.cpp index 0da306fcf..8aed378de 100644 --- a/common/CUDAWrapper.cpp +++ b/common/CUDAWrapper.cpp @@ -105,6 +105,10 @@ size_t get_gpu_total_memory(int32_t) { void set_gpu(int32_t dev_id) {} +bool gpu_direct_host_read() { + return false; +} + void pin_gpu() {} void pin_gpu(int32_t dev_id) {} diff --git a/common/CUDAWrapper.cu b/common/CUDAWrapper.cu index 4df5a9daa..ed7976a02 100644 --- a/common/CUDAWrapper.cu +++ b/common/CUDAWrapper.cu @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only +#include #include #include #include @@ -66,6 +67,28 @@ size_t get_gpu_total_memory(int32_t dev_id) { return prop.totalGlobalMem; } +bool gpu_direct_host_read() { + // Pageable memory access through the host's page tables is the hardware-coherent case (ATS over + // NVLink-C2C); without the second attribute the access is HMM-emulated, page faults and migration, + // and slower than a copy. + static const std::vector direct = [] { + const bool force = std::getenv("JFJOCH_FORCE_DIRECT_HOST_READ") != nullptr; + std::vector ret(get_gpu_count(), force); + for (int d = 0; d < static_cast(ret.size()); d++) { + int pageable = 0, host_page_tables = 0; + cudaDeviceGetAttribute(&pageable, cudaDevAttrPageableMemoryAccess, d); + cudaDeviceGetAttribute(&host_page_tables, cudaDevAttrPageableMemoryAccessUsesHostPageTables, d); + if (pageable && host_page_tables) + ret[d] = true; + } + return ret; + }(); + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess || device >= static_cast(direct.size())) + return false; + return direct[device]; +} + void set_gpu(int32_t dev_id) { auto dev_count = get_gpu_count(); diff --git a/common/CUDAWrapper.h b/common/CUDAWrapper.h index 76f5cf23a..d89d93053 100644 --- a/common/CUDAWrapper.h +++ b/common/CUDAWrapper.h @@ -31,6 +31,14 @@ size_t get_gpu_total_memory(int32_t dev_id); void set_gpu(int32_t dev_id); +// Whether kernels on the current device read ordinary (pageable) host memory directly and at memory +// speed, so a host buffer can be handed to a kernel instead of being copied to the card first: true on +// a unified-memory system where the GPU walks the CPU's page tables (DGX Spark GB10, Grace Hopper GH200, +// GB200), false on a discrete card. Asked of the driver once per device. JFJOCH_FORCE_DIRECT_HOST_READ +// in the environment makes it true on every device - for testing the direct path on a discrete card, +// where only page-locked host memory is then readable. false without CUDA. +bool gpu_direct_host_read(); + // Pin the calling thread to the next GPU in round-robin order, using a process-wide counter // (counter++ % get_gpu_count()). Call once per thread; no thread id needed. No-op when no GPU // is visible. Honours CUDA_VISIBLE_DEVICES via get_gpu_count(). diff --git a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu index cbe760fcd..ab5bf368b 100644 --- a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu +++ b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu @@ -7,6 +7,7 @@ #include "BSLZ4DecoderGPU.h" #include "BSLZ4DecodeWarp.h" +#include "../../common/CUDAWrapper.h" #include "../../common/JFJochException.h" #include "../../compression/JFJochDecompress.h" // BSHUF_BLOCKED_MULT and the container layout @@ -118,6 +119,7 @@ bool BSLZ4DecoderGPU::Supports(const CompressedImage &image) { BSLZ4DecoderGPU::BSLZ4DecoderGPU(size_t in_max_uncompressed_bytes, std::shared_ptr in_stream) : stream(std::move(in_stream)), + direct_host_read(gpu_direct_host_read()), max_uncompressed_bytes(in_max_uncompressed_bytes) { gpu_status = CudaDevicePtr(1); host_status = CudaHostPtr(1); @@ -211,7 +213,8 @@ BSLZ4ShuffledImage BSLZ4DecoderGPU::PrepareChunk(const CompressedImage &image) { if (nblocks_needed > (clen - 12) / MIN_BLOCK_BYTES_ON_WIRE) throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk too short for the blocks its header implies"); - EnsureCompressedCapacity(clen); + if (!direct_host_read) + EnsureCompressedCapacity(clen); EnsureBlockCapacity(nblocks_needed); size_t nblk = 0, off = 12, out_off = 0; @@ -240,8 +243,14 @@ BSLZ4ShuffledImage BSLZ4DecoderGPU::PrepareChunk(const CompressedImage &image) { host_status.get()[0] = 0; cuda_err(cudaMemcpyAsync(gpu_status.get(), host_status.get(), sizeof(uint32_t), cudaMemcpyHostToDevice, *stream)); - memcpy(host_compressed.get(), src, clen); - cuda_err(cudaMemcpyAsync(gpu_compressed.get(), host_compressed.get(), clen, cudaMemcpyHostToDevice, *stream)); + // Where the card reads host memory directly, the kernels read the reader's own buffer, which the + // caller keeps until it has synchronised the stream. + const uint8_t *compressed = src; + if (!direct_host_read) { + memcpy(host_compressed.get(), src, clen); + cuda_err(cudaMemcpyAsync(gpu_compressed.get(), host_compressed.get(), clen, cudaMemcpyHostToDevice, *stream)); + compressed = gpu_compressed.get(); + } const int nb = static_cast(nblk); if (nb > 0) @@ -249,7 +258,7 @@ BSLZ4ShuffledImage BSLZ4DecoderGPU::PrepareChunk(const CompressedImage &image) { cudaMemcpyHostToDevice, *stream)); BSLZ4ShuffledImage ret; - ret.compressed = gpu_compressed.get(); + ret.compressed = compressed; ret.desc = gpu_desc.get(); ret.status = gpu_status.get(); ret.nblocks = nb; @@ -257,7 +266,7 @@ BSLZ4ShuffledImage BSLZ4DecoderGPU::PrepareChunk(const CompressedImage &image) { ret.block_bytes = block_bytes; ret.tail_elems = static_cast(leftover_bytes / elem_size); ret.tail_elem0 = static_cast(out_off / elem_size); - ret.tail_src = leftover_bytes > 0 ? gpu_compressed.get() + off : nullptr; + ret.tail_src = leftover_bytes > 0 ? compressed + off : nullptr; return ret; } @@ -306,11 +315,12 @@ void BSLZ4DecoderGPU::Decode(const CompressedImage &image, uint8_t *gpu_out) { cuda_err(cudaGetLastError()); } - // The verbatim tail is already on the device inside the uploaded chunk. + // The verbatim tail is already on the device inside the uploaded chunk - or in host memory the card + // reads directly, hence cudaMemcpyDefault. if (s.tail_elems > 0) cuda_err(cudaMemcpyAsync(gpu_out + static_cast(s.tail_elem0) * s.elem_size, s.tail_src, static_cast(s.tail_elems) * s.elem_size, - cudaMemcpyDeviceToDevice, *stream)); + cudaMemcpyDefault, *stream)); } void BSLZ4DecoderGPU::ThrowIfDecodeFailed() { diff --git a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h index f6c306a98..de14c859d 100644 --- a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h +++ b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h @@ -66,6 +66,9 @@ class BSLZ4DecoderGPU { // uploads pageable memory through small staging buffers of its own, a piece at a time - at a few // GB/s on a host whose memory is far from the card. Copied here, the upload is one DMA at the link rate. CudaHostPtr host_compressed; + // The card reads the chunk where the reader left it (gpu_direct_host_read()); neither buffer above + // is then used. + bool direct_host_read = false; CudaDevicePtr gpu_shuffled; // LZ4 output, still bitshuffled CudaDevicePtr gpu_desc; CudaHostPtr host_desc; // pinned, so the descriptor upload is truly async diff --git a/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu b/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu index 41362f129..273afbc23 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu +++ b/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu @@ -7,6 +7,7 @@ #include "ImagePreprocessorGPU.h" #include "BSLZ4DecodeWarp.h" +#include "../../common/CUDAWrapper.h" #include "../../common/JFJochException.h" namespace { @@ -14,6 +15,12 @@ namespace { if (val != cudaSuccess) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); } + + // A kernel reads pixels of type T, so a host buffer it reads in place must be aligned to T. A + // decompressed buffer always is; an image inside a received message need not be. + bool aligned_to(const void *ptr, size_t bytes) { + return reinterpret_cast(ptr) % bytes == 0; + } } template @@ -338,6 +345,7 @@ ImagePreprocessorGPU::ImagePreprocessorGPU(const DiffractionExperiment &experime threads = 128; blocks = 4 * prop.multiProcessorCount; + direct_host_read = gpu_direct_host_read(); } float ImagePreprocessorGPU::GetLastDecompressionTime_s() const { @@ -359,7 +367,7 @@ ImagePreprocessorGPU::~ImagePreprocessorGPU() { } void ImagePreprocessorGPU::PinInputRegion(const void *ptr, size_t bytes) { - if (ptr == pinned_input && bytes == pinned_input_bytes) + if (direct_host_read || (ptr == pinned_input && bytes == pinned_input_bytes)) return; if (pinned_input) cuda_err(cudaHostUnregister(const_cast(pinned_input))); @@ -385,7 +393,8 @@ void ImagePreprocessorGPU::PinInputBuffer(std::vector &buffer, size_t s // Unregister before the resize, which can move the buffer. input_reg.unregister(); buffer.resize(size); - input_reg.rebind(buffer); + if (!direct_host_read) + input_reg.rebind(buffer); } ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *image_ptr, @@ -527,6 +536,11 @@ ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed const uint8_t *input, T err_value, T sat_value) { + // Read in place where the card reads host memory directly; the input outlives this call, which + // waits for the stream. + if (direct_host_read && aligned_to(input, sizeof(T))) + return AnalyzeOnDevice(processed_image, input, err_value, sat_value); + // Allocated here rather than in the constructor: only the host-upload path needs it, and the // device-decode path - which is what BSHUF_LZ4 images take - never touches it. Overshoot to // 4 bytes per pixel so a 1- or 2-byte image fits the same buffer. @@ -537,14 +551,11 @@ ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed // every blocking stream in the process, which serialised all workers behind whichever one was // uploading. The stream is synchronised at the end of this function, so the ordering is unchanged. cuda_err(cudaMemcpyAsync(gpu_decompressed_image, input, npixels * sizeof(T), cudaMemcpyHostToDevice, *stream)); - return AnalyzeOnDevice(processed_image, err_value, sat_value); + return AnalyzeOnDevice(processed_image, gpu_decompressed_image.get(), err_value, sat_value); } ImageStatistics ImagePreprocessorGPU::AnalyzeStoredInt32(ImagePreprocessorBuffer &processed_image, const CompressedImage &image) { - if (!gpu_decompressed_image.get()) - gpu_decompressed_image = CudaDevicePtr(npixels * sizeof(uint32_t)); - if (image.GetCompressionAlgorithm() == CompressionAlgorithm::UINT16_WIDENED) { // Uploaded as stored, two bytes a pixel, and widened by the preprocessing kernel itself: with // int32's error marker and int32's saturation, each pixel is decided exactly as the int32 image @@ -552,31 +563,40 @@ ImageStatistics ImagePreprocessorGPU::AnalyzeStoredInt32(ImagePreprocessorBuffer last_decode = LastDecode::None; if (image.GetCompressedSize() != npixels * sizeof(uint16_t)) throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in size"); + // Read in place where the card reads host memory directly: nothing is page-locked, so engines + // sharing the frame do not interfere. + if (direct_host_read && aligned_to(image.GetCompressed(), sizeof(uint16_t))) + return AnalyzeOnDevice(processed_image, image.GetCompressed(), INT32_MIN, INT32_MAX); // Staged through this engine's own page-locked buffer, as the CBF and BSLZ4 decoders stage // theirs, and not page-locked in place: the pre-scan hands one kept frame to several engines // at once, and an engine that unregisters a frame another one is still copying from takes the // whole context down. AnalyzeOnDevice waits for the stream, so the buffer is free again here. + if (!gpu_decompressed_image.get()) + gpu_decompressed_image = CudaDevicePtr(npixels * sizeof(uint32_t)); if (!host_stored.get()) host_stored = CudaHostPtr(npixels * sizeof(uint16_t)); std::memcpy(host_stored.get(), image.GetCompressed(), image.GetCompressedSize()); cuda_err(cudaMemcpyAsync(gpu_decompressed_image, host_stored.get(), image.GetCompressedSize(), cudaMemcpyHostToDevice, *stream)); - return AnalyzeOnDevice(processed_image, INT32_MIN, INT32_MAX); + return AnalyzeOnDevice(processed_image, gpu_decompressed_image.get(), INT32_MIN, INT32_MAX); } + if (!gpu_decompressed_image.get()) + gpu_decompressed_image = CudaDevicePtr(npixels * sizeof(uint32_t)); last_decode = LastDecode::CBF; if (!cbf_decoder) cbf_decoder = std::make_unique(stream); cbf_decoder->Decode(image, reinterpret_cast(gpu_decompressed_image.get())); - const ImageStatistics ret = AnalyzeOnDevice(processed_image, INT32_MIN, INT32_MAX); + const ImageStatistics ret = AnalyzeOnDevice(processed_image, gpu_decompressed_image.get(), + INT32_MIN, INT32_MAX); cbf_decoder->ThrowIfDecodeFailed(); return ret; } -// Everything after the image is on the device, shared by the host-upload and the device-decode -// entry points so the two cannot drift apart. +// Everything after the image is on the device - or in host memory the card reads directly - shared by +// the host-upload and the device-decode entry points so the two cannot drift apart. template -ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image, +ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image, const void *input, int64_t err_value, int64_t sat_value) { if (sat_value > saturation_limit) sat_value = saturation_limit; @@ -584,7 +604,7 @@ ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &p cpu_stats[0] = ImageStatistics{.max_value = INT64_MIN, .min_value = INT64_MAX}; cuda_err(cudaMemcpyAsync(gpu_stats, cpu_stats.data(), sizeof(ImageStatistics), cudaMemcpyHostToDevice, *stream)); preprocess_kernel <<< blocks, threads, 0, *stream >>>( - reinterpret_cast(gpu_decompressed_image.get()), + reinterpret_cast(input), gpu_mask->get(), processed_image.getGPUBuffer(), gpu_stats, diff --git a/image_analysis/image_preprocessing/ImagePreprocessorGPU.h b/image_analysis/image_preprocessing/ImagePreprocessorGPU.h index 7b2b60c84..df235f0eb 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorGPU.h +++ b/image_analysis/image_preprocessing/ImagePreprocessorGPU.h @@ -31,6 +31,9 @@ class ImagePreprocessorGPU : public ImagePreprocessor { size_t pinned_input_bytes = 0; // Staging for a marCCD frame's stored 16-bit pixels (UINT16_WIDENED), allocated on first use. CudaHostPtr host_stored; + // The card reads host memory directly (gpu_direct_host_read()): images are read where they lie, + // nothing is page-locked or staged; only the CBF device decoder still uploads its input. + bool direct_host_read = false; std::vector cpu_image; @@ -46,10 +49,11 @@ class ImagePreprocessorGPU : public ImagePreprocessor { ImageStatistics AnalyzeStoredInt32(ImagePreprocessorBuffer &processed_image, const CompressedImage &image); template ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *input, T err_value, T sat_value); - // Preprocess an image already sitting in gpu_decompressed_image (the host-upload path). - // err_value and sat_value are compared with the pixel's value, not its bits, so a narrower T can - // carry the int32 semantics of the image it widens to. - template ImageStatistics AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image, + // Preprocess the image at `input`: gpu_decompressed_image on the host-upload path, or the host + // buffer itself where the card reads host memory directly. err_value and sat_value are compared + // with the pixel's value, not its bits, so a narrower T can carry the int32 semantics of the image + // it widens to. + template ImageStatistics AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image, const void *input, int64_t err_value, int64_t sat_value); // Preprocess straight out of the bitshuffled bytes (the device-decode path). Same per-pixel // decision and same statistics as AnalyzeOnDevice, with the un-transpose folded in. diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 2582589b5..60ff860cb 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -2789,11 +2789,17 @@ void RotationScaleMergeGPU::IngestUpload(int n_frames, const Reflection *const * CudaCheck(cudaGetLastError(), "IngestFrameKernel launch"); // Through two pinned buffers of whole frames, each filled on all threads while the other is on - // its way to the card and taken apart there. + // its way to the card and taken apart there. Where the card reads host memory directly + // (gpu_direct_host_read()), the kernel takes the pinned buffer apart in place, without the upload. constexpr size_t STAGE_BYTES = 64 << 20; + const bool direct = gpu_direct_host_read(); const int cap = std::max(static_cast(STAGE_BYTES / sizeof(Reflection)), max_count); CudaHostPtr stage[2] = {CudaHostPtr(cap), CudaHostPtr(cap)}; - CudaDevicePtr staged[2] = {d.Alloc(cap), d.Alloc(cap)}; + CudaDevicePtr staged[2]; + if (!direct) { + staged[0] = d.Alloc(cap); + staged[1] = d.Alloc(cap); + } CudaEvent copied[2]; bool in_flight[2] = {false, false}; IngestArrays a{d.obs.get(), d.zeta.get(), d.corr.get(), d.bkg.get(), d.var_bkg.get(), d.image_number.get(), @@ -2814,12 +2820,19 @@ void RotationScaleMergeGPU::IngestUpload(int n_frames, const Reflection *const * std::copy(refl[o], refl[o] + count[o], dst + (d.in_src_start[o] - base)); }); if (cnt > 0) { - CudaCheck(cudaMemcpyAsync(staged[b].get(), dst, size_t(cnt) * sizeof(Reflection), - cudaMemcpyHostToDevice, d.s()), "upload reflections"); - CudaCheck(cudaEventRecord(copied[b], d.s()), "upload reflections"); + const Reflection *src = dst; + if (!direct) { + CudaCheck(cudaMemcpyAsync(staged[b].get(), dst, size_t(cnt) * sizeof(Reflection), + cudaMemcpyHostToDevice, d.s()), "upload reflections"); + CudaCheck(cudaEventRecord(copied[b], d.s()), "upload reflections"); + src = staged[b].get(); + } in_flight[b] = true; - IngestExtractKernel<<<(cnt + BLK - 1) / BLK, BLK, 0, d.s()>>>(cnt, base, staged[b].get(), a); + IngestExtractKernel<<<(cnt + BLK - 1) / BLK, BLK, 0, d.s()>>>(cnt, base, src, a); CudaCheck(cudaGetLastError(), "IngestExtractKernel launch"); + // Read in place, the buffer is free again only once the kernel has run. + if (direct) + CudaCheck(cudaEventRecord(copied[b], d.s()), "extract reflections"); } b ^= 1; f0 = f1; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 6f4632394..6f102faa2 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -129,6 +129,7 @@ ADD_EXECUTABLE(jfjoch_test BSLZ4DecoderGPUFuzzTest.cpp ImagePreprocessorGPUFusedTest.cpp StoredImageDecodeGPUTest.cpp + DirectHostReadGPUTest.cpp AdaptiveSpotFinderGPUTest.cpp SpotExtractorGPUParityTest.cpp CUDAMemHelpersTest.cpp diff --git a/tests/DirectHostReadGPUTest.cpp b/tests/DirectHostReadGPUTest.cpp new file mode 100644 index 000000000..2a150aac7 --- /dev/null +++ b/tests/DirectHostReadGPUTest.cpp @@ -0,0 +1,148 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include +#include "../common/CUDAWrapper.h" + +#ifdef JFJOCH_USE_CUDA + +#include +#include +#include + +#include "../common/PixelMask.h" +#include "../compression/JFJochCompressor.h" +#include "../image_analysis/image_preprocessing/BSLZ4DecoderGPU.h" +#include "../image_analysis/image_preprocessing/ImagePreprocessorCPU.h" +#include "../image_analysis/image_preprocessing/ImagePreprocessorGPU.h" +#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h" + +// The images the device reads straight out of host memory where gpu_direct_host_read() holds (a +// unified-memory system), against the host decoder and the CPU preprocessor. The host buffers are +// page-locked, so a discrete card can read them too: run with JFJOCH_FORCE_DIRECT_HOST_READ=1 to take +// the direct path there; without it these check the copying path on the same inputs. +namespace { + +bool SameStats(const ImageStatistics &a, const ImageStatistics &b) { + return a.max_value == b.max_value && a.min_value == b.min_value + && a.masked_pixel_count == b.masked_pixel_count + && a.error_pixel_count == b.error_pixel_count + && a.saturated_pixel_count == b.saturated_pixel_count; +} + +std::vector MakeImage(size_t npixels, uint32_t seed) { + std::mt19937 rng(seed); + std::vector img(npixels, 0); + for (size_t i = npixels / 4; i < npixels / 2; i++) + img[i] = rng() % 11; + for (size_t s = 0; s < 500; s++) + img[rng() % npixels] = 30000 + rng() % 50000; + for (size_t s = 0; s < 300; s++) + img[rng() % npixels] = UINT32_MAX; + return img; +} + +void CheckSameAsCPU(const DiffractionExperiment &x, const PixelMask &mask, const ImagePreprocessorBufferGPU &gpu_buf, + const ImageStatistics &gpu_stats, const uint8_t *raw, CompressedImageMode mode) { + const size_t npixels = x.GetPixelsNum(); + ImagePreprocessorCPU cpu_pre(x, mask); + ImagePreprocessorBuffer cpu_buf(npixels); + const ImageStatistics cpu_stats = cpu_pre.Analyze(cpu_buf, raw, mode); + CHECK(SameStats(cpu_stats, gpu_stats)); + size_t ndiff = 0; + for (size_t i = 0; i < npixels; i++) + ndiff += gpu_buf[i] != cpu_buf[i]; + CHECK(ndiff == 0); +} + +} // namespace + +TEST_CASE("DirectHostRead_BSLZ4Decode", "[direct_host_read]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + WARN("gpu_direct_host_read() = " << gpu_direct_host_read()); + + // An odd pixel count, so the chunk ends in a verbatim tail that is copied rather than decoded. + const size_t width = 1001, height = 997, npixels = width * height; + const auto original = MakeImage(npixels, 3); + JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4); + std::vector compressed = compressor.Compress(original); + CudaRegisteredVector reg(compressed); + const CompressedImage image(compressed.data(), compressed.size(), width, height, CompressedImageMode::Uint32, + CompressionAlgorithm::BSHUF_LZ4); + + auto stream = std::make_shared(); + BSLZ4DecoderGPU decoder(npixels * sizeof(uint32_t), stream); + CudaDevicePtr gpu_out(npixels * sizeof(uint32_t)); + decoder.Decode(image, gpu_out.get()); + REQUIRE(cudaStreamSynchronize(*stream) == cudaSuccess); + decoder.ThrowIfDecodeFailed(); + + std::vector result(npixels); + REQUIRE(cudaMemcpy(result.data(), gpu_out.get(), npixels * sizeof(uint32_t), cudaMemcpyDeviceToHost) + == cudaSuccess); + CHECK(result == original); +} + +TEST_CASE("DirectHostRead_Preprocessor", "[direct_host_read]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + + DetectorSetup detector = DetJF4M(); + detector.SaturationLimit(65535); + DiffractionExperiment x(detector); + x.DetectorDistance_mm(80).BeamX_pxl(1030).BeamY_pxl(1080); + const size_t npixels = x.GetPixelsNum(); + PixelMask mask(x); + std::vector user_mask(mask.GetMask().size(), 0); + for (size_t i = 0; i < npixels; i += 997) user_mask[i] = 1; + mask.LoadUserMask(x, user_mask); + + auto stream = std::make_shared(); + ImagePreprocessorGPU gpu_pre(x, mask, stream, /*copy_image_to_host=*/true); + + // The bitshuffle/LZ4 chunk, decoded and preprocessed by the fused kernel. + std::vector img = MakeImage(npixels, 5); + JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4); + std::vector compressed = compressor.Compress(img); + CudaRegisteredVector compressed_reg(compressed); + { + const CompressedImage image(compressed.data(), compressed.size(), x.GetXPixelsNum(), x.GetYPixelsNum(), + CompressedImageMode::Uint32, CompressionAlgorithm::BSHUF_LZ4); + ImagePreprocessorBufferGPU gpu_buf(npixels); + ImageStatistics gpu_stats{}; + REQUIRE(gpu_pre.AnalyzeCompressed(gpu_buf, image, gpu_stats)); + CheckSameAsCPU(x, mask, gpu_buf, gpu_stats, reinterpret_cast(img.data()), + CompressedImageMode::Uint32); + } + + // An uncompressed image, as the host-upload path receives it. + { + std::vector raw(reinterpret_cast(img.data()), + reinterpret_cast(img.data()) + npixels * sizeof(uint32_t)); + CudaRegisteredVector raw_reg(raw); + ImagePreprocessorBufferGPU gpu_buf(npixels); + const ImageStatistics gpu_stats = gpu_pre.Analyze(gpu_buf, raw.data(), CompressedImageMode::Uint32); + CheckSameAsCPU(x, mask, gpu_buf, gpu_stats, raw.data(), CompressedImageMode::Uint32); + } + + // A marCCD frame's stored 16-bit pixels, widened by the kernel. + { + std::mt19937 rng(7); + std::vector stored(npixels * sizeof(uint16_t)); + for (size_t i = 0; i < npixels; i++) { + const uint16_t v = (i % 1500 == 0) ? 65535 : static_cast(rng() % 70000); + std::memcpy(stored.data() + i * sizeof(uint16_t), &v, sizeof(v)); + } + CudaRegisteredVector stored_reg(stored); + const CompressedImage image(stored.data(), stored.size(), x.GetXPixelsNum(), x.GetYPixelsNum(), + CompressedImageMode::Int32, CompressionAlgorithm::UINT16_WIDENED); + ImagePreprocessorBufferGPU gpu_buf(npixels); + ImageStatistics gpu_stats{}; + REQUIRE(gpu_pre.AnalyzeCompressed(gpu_buf, image, gpu_stats)); + std::vector decoded; + CheckSameAsCPU(x, mask, gpu_buf, gpu_stats, image.GetUncompressedPtr(decoded), CompressedImageMode::Int32); + } +} + +#endif