Keep 16-bit images 16-bit through the GPU pipeline

A detector reading out 16 bits had its frame widened to int32 the moment it was
decoded, and every per-pixel pass over that frame then moved four bytes a pixel to
carry two. Those passes - the ring statistics three times over, the strong-pixel
search, spot extraction, the azimuthal and ROI integrators, Bragg integration - are
the bulk of the image loop's device traffic, and 16 bits is the mode a fast
acquisition runs in, which is exactly where throughput matters.

The preprocessed image now keeps the width of its source. Two codes at the top of the
16-bit range carry the two special states, and they cannot collide with a real value:

  0xFFFF          masked, or the source's own bad-pixel marker.
  saturation      a pixel at or above the saturation limit. 0xFFFE where the limit
  code            leaves room - a 16-bit EIGER declares a count-rate limit of a few
                  thousand, so there is room to spare - and 0xFFFF where the limit is
                  the whole range, in which case the "is error" test has already
                  claimed 0xFFFF, nothing can be saturated, and 0xFFFE stays a real
                  value.

Either way a real value is strictly below the saturation limit and so below both
codes. Nothing is clipped and nothing is lost, and which code is in force is carried
with the image rather than assumed.

No engine learns a second convention. PixelView widens on load, so a masked pixel
still reads as INT32_MIN and a saturated one as INT32_MAX, and every existing
`v != INT32_MIN && v != INT32_MAX` test keeps its meaning. One code path, not two
instantiations that can drift apart; the branch is on a pointer that is the same for
every thread of every block, on kernels whose time is the loads it selects between.
The vector loads are kept - four pixels still arrive in one transaction, 16 bytes wide
or 8, whichever the image is.

The wide path is unchanged, and is still taken for anything that is not a 16-bit
source, and for any caller that wants the preprocessed image copied back to the host -
that mirror is int32 and the CPU engines know only that convention.

Measured on the one 16-bit dataset in the rotation test set, which is also the
smallest detector in it (2.5M pixels, where per-pixel work is a small part of the
loop): image loop 1.025 s -> 1.005 s at one GPU, whole run 5.64 s -> 5.52 s. The gain
scales with the frame, so a 16M-pixel detector - where six full-frame passes are 86 %
of the loop's GPU time - has much more to gain, and nothing here can measure that:
every other dataset in the test set is stored 32-bit.

Correctness on that dataset is exact where it can be: indexing rate, first-pass
validation score and the integrated partial count are identical to the wide path, and
its whole battery row - reflections, observations, space group, R_meas, CC1/2, ISa,
mosaicity - is unchanged. Battery 6m17s -> 6m16s, 21/24 space groups, no failures.

Also logs, once per run, the width the images are stored in, since it decides how much
of the frame moves through every pass.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

image_analysis/image_preprocessing/ImagePreprocessorGPU.cu

@@ -4,6 +4,7 @@
 #include <type_traits>
 
 #include "ImagePreprocessorGPU.h"
+#include "PreprocessedPixel.h"
 #include "../../common/JFJochException.h"
 
 namespace {
@@ -13,15 +14,46 @@ namespace {
     }
 }
 
-template<class T>
+
+// Store a finished pixel in the width this image uses. The narrow form is only ever chosen when the
+// source cannot produce a saturated pixel (see ChooseNarrow), so INT32_MIN - masked or bad - is the
+// only special value that can arrive here; it becomes the one reserved 16-bit code. INT32_MAX would
+// truncate to that same code, so even a broken invariant degrades to "bad" rather than to a wrong
+// intensity.
+template<class OutT> __device__ __forceinline__ OutT StorePixel(int32_t v, uint16_t sat_code);
+template<> __device__ __forceinline__ int32_t StorePixel<int32_t>(int32_t v, uint16_t) { return v; }
+template<> __device__ __forceinline__ uint16_t StorePixel<uint16_t>(int32_t v, uint16_t sat_code) {
+    if (v == INT32_MIN) return preprocessed_pixel::NARROW_BAD;
+    if (v == INT32_MAX) return sat_code;
+    return static_cast<uint16_t>(v);
+}
+
+
+// Eight finished pixels in one aligned store, in whichever width the image is stored in.
+template<class OutT> __device__ __forceinline__ void StorePixels8(OutT *dst, const int32_t o[8], uint16_t sat_code);
+template<> __device__ __forceinline__ void StorePixels8<int32_t>(int32_t *dst, const int32_t o[8], uint16_t) {
+    int4 *d = reinterpret_cast<int4 *>(dst);
+    d[0] = make_int4(o[0], o[1], o[2], o[3]);
+    d[1] = make_int4(o[4], o[5], o[6], o[7]);
+}
+template<> __device__ __forceinline__ void StorePixels8<uint16_t>(uint16_t *dst, const int32_t o[8], uint16_t sc) {
+    ushort4 *d = reinterpret_cast<ushort4 *>(dst);
+    d[0] = make_ushort4(StorePixel<uint16_t>(o[0], sc), StorePixel<uint16_t>(o[1], sc),
+                        StorePixel<uint16_t>(o[2], sc), StorePixel<uint16_t>(o[3], sc));
+    d[1] = make_ushort4(StorePixel<uint16_t>(o[4], sc), StorePixel<uint16_t>(o[5], sc),
+                        StorePixel<uint16_t>(o[6], sc), StorePixel<uint16_t>(o[7], sc));
+}
+
+template<class T, class OutT>
 __global__ void preprocess_kernel(
     const T *__restrict__ input,
     const uint8_t *__restrict__ mask,
-    int32_t *__restrict__ output,
+    OutT *__restrict__ output,
     ImageStatistics *__restrict__ stats,
     T saturation_limit,
     T err_value,
-    int npixels) {
+    int npixels,
+    uint16_t narrow_sat_code) {
     // Shared block accumulators
     __shared__ unsigned long long s_masked;
     __shared__ unsigned long long s_saturated;
@@ -60,8 +92,9 @@ __global__ void preprocess_kernel(
         bool valid = !(is_masked || is_sat || is_err);
 
         // Output
-        output[i] =
-                is_masked ? INT32_MIN : is_err ? INT32_MIN : is_sat ? INT32_MAX : (int32_t) v;
+        output[i] = StorePixel<OutT>(
+                is_masked ? INT32_MIN : is_err ? INT32_MIN : is_sat ? INT32_MAX : (int32_t) v,
+                narrow_sat_code);
 
         // Counters
         local_masked += is_masked;
@@ -166,15 +199,16 @@ __device__ __forceinline__ uint64_t transpose8_fused(uint64_t x) {
 //
 // The last CUDA block (blockIdx.x == nblocks) finishes the handful of elements bitshuffle stores
 // verbatim; they are already on the device inside the uploaded chunk.
-template<class T, int ES>
+template<class T, int ES, class OutT>
 __global__ __launch_bounds__(256) void untranspose_preprocess_kernel(
         const uint8_t *__restrict__ shuffled,
         const BSLZ4BlockDesc *__restrict__ desc,
         const uint8_t *__restrict__ mask,
-        int32_t *__restrict__ out,
+        OutT *__restrict__ out,
         ImageStatistics *__restrict__ stats,
         T sat_value, T err_value, int nblocks,
-        const uint8_t *__restrict__ tail_src, uint32_t tail_elems, uint32_t tail_elem0) {
+        const uint8_t *__restrict__ tail_src, uint32_t tail_elems, uint32_t tail_elem0,
+        uint16_t narrow_sat_code) {
     PreprocessAccum<T> l;
 
     // The bytes are assembled in the unsigned counterpart of T - shifting a byte into the top of a
@@ -189,7 +223,9 @@ __global__ __launch_bounds__(256) void untranspose_preprocess_kernel(
             for (int p = 0; p < ES; p++)
                 uv |= (U)((U)tail_src[threadIdx.x * ES + p] << (8 * p));
             const T v = (T) uv;
-            out[tail_elem0 + threadIdx.x] = l.Apply(v, mask[tail_elem0 + threadIdx.x] != 0, sat_value, err_value);
+            out[tail_elem0 + threadIdx.x] = StorePixel<OutT>(
+                    l.Apply(v, mask[tail_elem0 + threadIdx.x] != 0, sat_value, err_value),
+                    narrow_sat_code);
         }
         FlushStats<T>(l, stats);
         return;
@@ -219,10 +255,9 @@ __global__ __launch_bounds__(256) void untranspose_preprocess_kernel(
             for (int p = 0; p < ES; p++) uv |= (U)((U)((x[p] >> (8 * k)) & 0xff) << (8 * p));
             o[k] = l.Apply((T) uv, mask[elem0 + i * 8 + k] != 0, sat_value, err_value);
         }
-        // elem0 is a multiple of 8 (bitshuffle blocks are), so this is 32-byte aligned.
-        int4 *dst = reinterpret_cast<int4 *>(out + elem0 + i * 8);
-        dst[0] = make_int4(o[0], o[1], o[2], o[3]);
-        dst[1] = make_int4(o[4], o[5], o[6], o[7]);
+        // elem0 is a multiple of 8 (bitshuffle blocks are), so eight pixels are one aligned store of
+        // 32 bytes wide or 16, whichever this image is.
+        StorePixels8<OutT>(out + elem0 + i * 8, o, narrow_sat_code);
     }
     FlushStats<T>(l, stats);
 }
@@ -266,6 +301,17 @@ void ImagePreprocessorGPU::PinInputBuffer(std::vector<uint8_t> &buffer, size_t s
     input_reg.rebind(buffer);
 }
 
+
+// Whether this image can keep its 16-bit width instead of being widened to int32.
+//
+// Whenever the source is 16-bit unsigned: the two reserved codes carry masked/bad and saturated,
+// and a real value is always below both (see PreprocessedPixel.h). The one exception is a caller
+// that wants the preprocessed image copied back to the host - that mirror is int32, and the CPU
+// engines reading it know only that convention.
+bool ImagePreprocessorGPU::ChooseNarrow(bool source_is_uint16) const {
+    return source_is_uint16 && !copy_image_to_host;
+}
+
 ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *image_ptr,
                                               CompressedImageMode image_mode) {
     switch (image_mode) {
@@ -334,20 +380,31 @@ ImageStatistics ImagePreprocessorGPU::UntransposeAndAnalyze(ImagePreprocessorBuf
     cpu_stats[0] = ImageStatistics{.max_value = INT64_MIN, .min_value = INT64_MAX};
     cuda_err(cudaMemcpyAsync(gpu_stats, cpu_stats.data(), sizeof(ImageStatistics), cudaMemcpyHostToDevice, *stream));
 
+    const bool narrow = ChooseNarrow(std::is_same<T, uint16_t>::value);
+    const uint16_t sat_code = preprocessed_pixel::NarrowSaturatedCode(
+            std::min<int64_t>(saturation_limit, static_cast<int64_t>(sat_value)));
+    processed_image.SetNarrow(narrow, sat_code);
+
     // One CUDA block per bitshuffle block, plus one for the verbatim tail when there is one.
     const int nb = shuffled.nblocks + (shuffled.tail_elems > 0 ? 1 : 0);
-    untranspose_preprocess_kernel<T, ES> <<< nb, 256, 0, *stream >>>(
+    if (narrow)
+        untranspose_preprocess_kernel<T, ES, uint16_t> <<< nb, 256, 0, *stream >>>(
         shuffled.shuffled,
         shuffled.desc,
         gpu_mask->get(),
-        processed_image.getGPUBuffer(),
+        reinterpret_cast<uint16_t *>(processed_image.getGPUBuffer()),
         gpu_stats,
         sat_value,
         err_value,
         shuffled.nblocks,
         shuffled.tail_src,
         shuffled.tail_elems,
-        shuffled.tail_elem0);
+        shuffled.tail_elem0, sat_code);
+    else
+        untranspose_preprocess_kernel<T, ES, int32_t> <<< nb, 256, 0, *stream >>>(
+            shuffled.shuffled, shuffled.desc, gpu_mask->get(),
+            processed_image.getGPUBuffer(), gpu_stats, sat_value, err_value,
+            shuffled.nblocks, shuffled.tail_src, shuffled.tail_elems, shuffled.tail_elem0, sat_code);
     cuda_err(cudaGetLastError());
 
     if (copy_image_to_host)
@@ -393,14 +450,23 @@ 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<T> <<< blocks, threads, 0, *stream >>>(
-        reinterpret_cast<const T *>(gpu_decompressed_image.get()),
-        gpu_mask->get(),
-        processed_image.getGPUBuffer(),
-        gpu_stats,
-        sat_value,
-        err_value,
-        npixels);
+
+    const bool narrow = ChooseNarrow(std::is_same<T, uint16_t>::value);
+    const uint16_t sat_code = preprocessed_pixel::NarrowSaturatedCode(
+            std::min<int64_t>(saturation_limit, static_cast<int64_t>(sat_value)));
+    processed_image.SetNarrow(narrow, sat_code);
+    // The narrow image is the same allocation written two bytes at a time, so there is nothing to
+    // allocate and the wide pointer is simply reinterpreted.
+    if (narrow)
+        preprocess_kernel<T, uint16_t> <<< blocks, threads, 0, *stream >>>(
+            reinterpret_cast<const T *>(gpu_decompressed_image.get()), gpu_mask->get(),
+            reinterpret_cast<uint16_t *>(processed_image.getGPUBuffer()),
+            gpu_stats, sat_value, err_value, npixels, sat_code);
+    else
+        preprocess_kernel<T, int32_t> <<< blocks, threads, 0, *stream >>>(
+            reinterpret_cast<const T *>(gpu_decompressed_image.get()), gpu_mask->get(),
+            processed_image.getGPUBuffer(),
+            gpu_stats, sat_value, err_value, npixels, sat_code);
     cuda_err(cudaGetLastError());
     // The preprocessed image is 4 bytes per pixel - by far the largest transfer here - and every GPU
     // engine reads it straight from the device buffer, so it only comes back when a CPU engine needs it.