v1.0.0-rc.175 (#85)

* Rugnux: Processes a collection live while jfjoch_broker writes it (`rugnux <master.h5> <broker URL>`).
* Rugnux: Processing starts on the first images and confirms its choices once all are in, so analysis can begin during collection.
* Rugnux: Faster processing, especially on servers with several GPUs and CPU sockets and on computers without a GPU (`--no-gpu`, macOS).
* Rugnux: Uses much less GPU memory, so large data sets stay on the GPU of smaller cards, and a data set too large for the card is processed on the CPU instead.
* Rugnux: On macOS, Bragg integration runs on the Apple GPU.
* Rugnux: Better handling of twinning and pseudo-symmetry, with the twin hypothesis written as `<prefix>_twin_<SG>.mtz`, and a wrong `--model` no longer confirms the data or decides the hand.
* Rugnux: `--model` accepts predicted (AlphaFold) models and falls back to molecular replacement when the model is not placed as a rigid body (experimental `--mr` forces it), settling space group, hand and twinning and writing a different answer as `<prefix>_mr.mtz`.
* Rugnux: Molecular replacement is faster and settles the space group in more cases with several copies in the asymmetric unit.
* Rugnux: Deposition-ready mmCIF, a richer report, and numbered steps on the terminal with the full log in `<prefix>_log.txt`; `_P1.mtz`, `--estimate-beam-center` and `--beam-center-search` are removed.
* Rugnux: The beam-stop shadow is found correctly on tilted detectors.
* Rugnux: Stills files that record a stationary goniometer axis get the default stills geometry refinement.
* jfjoch_viewer: "Analyze dataset" uses the same defaults as a bare `rugnux <file>`, and shows the results of processed CBF, marCCD and SMV sweeps.
* jfjoch_viewer: The Files panel refreshes on F5 and every few seconds, and the plot type and other selections stay put during live collection.
* jfjoch_writer: Memory stays bounded and old data files are closed efficiently on long runs, any failure ends the run in an error state the broker is told about, and the master file appears only after its data files; the ZeroMQ pusher logs when a stalled writer makes it drop images.
* Fixes: The TCP image pusher no longer blocks the broker's receiver, GPU out-of-memory failures stop with a clear message, and the XDS plugin masks module and chip edges.

---------

Co-authored-by: Filip Leonarski <jungfrau@filips-mac-mini.home>
Reviewed-on: #85
Co-authored-by: Filip Leonarski <filip.leonarski@psi.ch>

image_analysis/image_preprocessing/ImagePreprocessorGPU.cu

@@ -2,10 +2,12 @@
 // SPDX-License-Identifier: GPL-3.0-only
 
 #include <algorithm>
+#include <cstring>
 #include <type_traits>
 
 #include "ImagePreprocessorGPU.h"
 #include "BSLZ4DecodeWarp.h"
+#include "../../common/CUDAWrapper.h"
 #include "../../common/JFJochException.h"
 
 namespace {
@@ -13,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<uintptr_t>(ptr) % bytes == 0;
+    }
 }
 
 template<class T>
@@ -21,8 +29,8 @@ __global__ void preprocess_kernel(
     const uint8_t *__restrict__ mask,
     int32_t *__restrict__ output,
     ImageStatistics *__restrict__ stats,
-    T saturation_limit,
-    T err_value,
+    int64_t saturation_limit,
+    int64_t err_value,
     int npixels) {
     // Shared block accumulators
     __shared__ unsigned long long s_masked;
@@ -262,8 +270,14 @@ __global__ __launch_bounds__(256) void untranspose_preprocess_kernel(
 // 16 kB blocks 32-bit detector data comes in, against the thirty-two warps the standalone kernel
 // keeps in flight. So this trades decode parallelism away for decode latency, and which way that
 // comes out is a measurement rather than an argument.
+//
+// Five resident blocks of 256 threads is what the shared buffer allows on an sm_120 card (100 kB per
+// SM), but the compiler's register choice held it to four; asking for five takes the kernel from
+// ~390 to ~315 us per 18 Mpx frame there. Where five cannot be resident anyway (sm_75: 1024 threads
+// per SM) the bound changes nothing. The block is never wider than 256 threads: the shared buffer is
+// at most FUSED_MAX_BLOCK_BYTES, and the launch gives it one thread per 64 bytes.
 template<class T, int ES>
-__global__ void decode_untranspose_preprocess_kernel(
+__global__ __launch_bounds__(256, 5) void decode_untranspose_preprocess_kernel(
         const uint8_t *__restrict__ src,
         const BSLZ4BlockDesc *__restrict__ desc,
         uint32_t *__restrict__ status,
@@ -314,8 +328,7 @@ ImagePreprocessorGPU::ImagePreprocessorGPU(const DiffractionExperiment &experime
       stream(stream),
       copy_image_to_host(copy_image_to_host),
       gpu_stats(1),
-      cpu_stats(1),
-      cpu_stats_reg(cpu_stats) {
+      cpu_stats(1) {
     // Setup mask. The byte-per-pixel form and its checksum come from the PixelMask, which derives them
     // whenever the mask changes: they are the same for every worker, and deriving them walks every
     // pixel of the detector - 18 million of them on a 16 Mpx one, per engine, with an engine built per
@@ -332,21 +345,29 @@ ImagePreprocessorGPU::ImagePreprocessorGPU(const DiffractionExperiment &experime
 
     threads = 128;
     blocks = 4 * prop.multiProcessorCount;
+    direct_host_read = gpu_direct_host_read();
 }
 
 float ImagePreprocessorGPU::GetLastDecompressionTime_s() const {
-    return bslz4_decoder ? bslz4_decoder->GetDecodeTime_s() : 0.0f;
+    switch (last_decode) {
+        case LastDecode::BSLZ4: return bslz4_decoder->GetDecodeTime_s();
+        case LastDecode::CBF: return cbf_decoder->GetDecodeTime_s();
+        default: return 0.0f;
+    }
 }
 
 ImagePreprocessorGPU::~ImagePreprocessorGPU() {
     // Unchecked, as everywhere else in this teardown path: a destructor is noexcept and the region
-    // may already be gone with the device.
+    // may already be gone with the device. The stream is waited for first: its pooled buffers are
+    // freed next (see CudaDevicePtr), and the region may still be the source of a copy.
+    if (stream)
+        cudaStreamSynchronize(*stream);
     if (pinned_input)
         cudaHostUnregister(const_cast<void *>(pinned_input));
 }
 
 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<void *>(pinned_input)));
@@ -372,7 +393,8 @@ void ImagePreprocessorGPU::PinInputBuffer(std::vector<uint8_t> &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,
@@ -398,8 +420,14 @@ ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed
 bool ImagePreprocessorGPU::AnalyzeCompressed(ImagePreprocessorBuffer &processed_image,
                                              const CompressedImage &image,
                                              ImageStatistics &stats) {
+    if (image.GetCompressionAlgorithm() == CompressionAlgorithm::CBF_BYTE_OFFSET
+        || image.GetCompressionAlgorithm() == CompressionAlgorithm::UINT16_WIDENED) {
+        stats = AnalyzeStoredInt32(processed_image, image);
+        return true;
+    }
     if (!BSLZ4DecoderGPU::Supports(image))
         return false;   // caller decompresses on the host and uses Analyze()
+    last_decode = LastDecode::BSLZ4;
     if (image.GetUncompressedSize() != npixels * image.GetByteDepth())
         return false;
 
@@ -508,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<T>(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.
@@ -518,21 +551,60 @@ 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<T>(processed_image, err_value, sat_value);
+    return AnalyzeOnDevice<T>(processed_image, gpu_decompressed_image.get(), err_value, sat_value);
 }
 
-// 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.
+ImageStatistics ImagePreprocessorGPU::AnalyzeStoredInt32(ImagePreprocessorBuffer &processed_image,
+                                                         const CompressedImage &image) {
+    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
+        // the host would have widened it to.
+        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<uint16_t>(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<uint8_t>(npixels * sizeof(uint32_t));
+        if (!host_stored.get())
+            host_stored = CudaHostPtr<uint8_t>(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<uint16_t>(processed_image, gpu_decompressed_image.get(), INT32_MIN, INT32_MAX);
+    }
+
+    if (!gpu_decompressed_image.get())
+        gpu_decompressed_image = CudaDevicePtr<uint8_t>(npixels * sizeof(uint32_t));
+    last_decode = LastDecode::CBF;
+    if (!cbf_decoder)
+        cbf_decoder = std::make_unique<CBFDecoderGPU>(stream);
+    cbf_decoder->Decode(image, reinterpret_cast<int32_t *>(gpu_decompressed_image.get()));
+    const ImageStatistics ret = AnalyzeOnDevice<int32_t>(processed_image, gpu_decompressed_image.get(),
+                                                         INT32_MIN, INT32_MAX);
+    cbf_decoder->ThrowIfDecodeFailed();
+    return ret;
+}
+
+// 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<class T>
-ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image,
-                                                      T err_value, T sat_value) {
+ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image, const void *input,
+                                                      int64_t err_value, int64_t sat_value) {
     if (sat_value > saturation_limit)
-        sat_value = static_cast<T>(saturation_limit);
+        sat_value = saturation_limit;
 
     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<T> <<< blocks, threads, 0, *stream >>>(
-        reinterpret_cast<const T *>(gpu_decompressed_image.get()),
+        reinterpret_cast<const T *>(input),
         gpu_mask->get(),
         processed_image.getGPUBuffer(),
         gpu_stats,