diff --git a/image_analysis/beam_stop/ShadowAccumulatorGPU.cu b/image_analysis/beam_stop/ShadowAccumulatorGPU.cu index efbb9130..be7117b0 100644 --- a/image_analysis/beam_stop/ShadowAccumulatorGPU.cu +++ b/image_analysis/beam_stop/ShadowAccumulatorGPU.cu @@ -105,16 +105,18 @@ bool ShadowAccumulatorGPU::Supports(const CompressedImage &image) { } void ShadowAccumulatorGPU::EnsureRawCapacity(size_t bytes_per_frame) { + if (bytes_per_frame * BATCH > raw_capacity) { + // The constructor already sized this for the widest pixel type, so in practice this only + // runs if that guess was too small. cudaMalloc and cudaFree both synchronise the whole + // device, which is why it is never done per frame. + FoldPending(); + cuda_err(cudaStreamSynchronize(*stream)); + raw_capacity = bytes_per_frame * BATCH; + raw = CudaDevicePtr(raw_capacity); + } + // Only once the pending batch has been folded: FoldPending strides the buffer by frame_bytes, + // so it has to keep describing the frames already in it until they are gone. frame_bytes = bytes_per_frame; - if (bytes_per_frame * BATCH <= raw_capacity) - return; - // The constructor already sized this for the widest pixel type, so in practice this only runs - // if that guess was too small. cudaMalloc and cudaFree both synchronise the whole device, which - // is why it is never done per frame. - FoldPending(); - cuda_err(cudaStreamSynchronize(*stream)); - raw_capacity = bytes_per_frame * BATCH; - raw = CudaDevicePtr(raw_capacity); } // Fold the frames decoded so far into the projection. One pass over the accumulator for the whole @@ -150,12 +152,15 @@ void ShadowAccumulatorGPU::Add(const CompressedImage &image) { throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "ShadowAccumulatorGPU: image size does not match the detector"); - if (!decoder) + // A batch holds one pixel type and one frame size; a change of either closes the batch first, + // while frame_bytes and pending_mode still describe the frames already in it. + if (pending > 0 && (image.GetMode() != pending_mode + || image.GetUncompressedSize() != frame_bytes)) + FoldPending(); + + if (!decoder || image.GetUncompressedSize() != frame_bytes) decoder = std::make_unique(image.GetUncompressedSize(), stream); EnsureRawCapacity(image.GetUncompressedSize()); - // A batch holds one pixel type; a change of depth mid-run closes the batch first. - if (pending > 0 && image.GetMode() != pending_mode) - FoldPending(); pending_mode = image.GetMode(); decoder->Decode(image, raw.get() + static_cast(pending) * frame_bytes); diff --git a/image_analysis/beam_stop/ShadowFinder.cpp b/image_analysis/beam_stop/ShadowFinder.cpp index 660831fc..2e1fe240 100644 --- a/image_analysis/beam_stop/ShadowFinder.cpp +++ b/image_analysis/beam_stop/ShadowFinder.cpp @@ -427,10 +427,9 @@ void ShadowFinder::AddImage(const DataMessage &data, std::vector &buffe "ShadowFinder: shard out of range"); #ifdef JFJOCH_USE_CUDA - if (ShadowAccumulatorGPU::Supports(data.image) && Gpu()) { - // One device, so the frames queue here - but each is only a chunk upload plus two kernels, - // and nothing was decompressed on the host to get this far. - ShadowAccumulatorGPU *acc = Gpu(); + // One device, so the frames queue here - but each is only a chunk upload plus two kernels, and + // nothing was decompressed on the host to get this far. + if (ShadowAccumulatorGPU *acc = ShadowAccumulatorGPU::Supports(data.image) ? Gpu() : nullptr) { std::unique_lock ul(gpu_mutex); try { acc->Add(data.image);