Give each FFT direction a block and its histogram shared memory

The de-novo indexer projects every spot onto each of 16384 search directions and
bins the projections; the peak of each direction's spectrum is a reciprocal
lattice row spacing. One thread owned a whole direction, so neighbouring lanes
wrote 12.6 kB apart and every warp instruction touched 32 separate sectors of a
206 MB buffer with no chance of staying in a 4 MB L2. 160 million scattered
global read-modify-writes, at about 14% of the card's bandwidth.

One block per direction now, with the bins in shared memory. They are counts, so
they are held as integers: an integer atomicAdd is a real shared-memory
instruction where the float one compiles to a compare-and-swap retry loop, and a
count below 2^24 converts to float exactly, so the output is bit for bit what the
repeated += 1.0 produced. Above 48 kB of bins the old kernel still runs.

245.75 ms per launch -> 4.08 ms, so 0.98 s of the run -> 0.016 s. This machine
runs two of them at once on two cards, so it is worth about half a second here
and about a second on the single-GPU machines the viewer and the broker run on.
The FFT it feeds takes 3.3 ms; preparing its input took 70x longer than
transforming it.

Alongside it, the per-frame scale fit divided by k^2 once per observation per
IRLS iteration, and a loop-invariant divisor does not get hoisted out of a double
division - ptxas emits the whole Newton refinement of the reciprocal every time.
Hoisted, as 1/sigma already is a few lines above; the same expression in the
three CPU scale paths went with it so the two stay algebraically identical.
54.92 ms per launch -> 44.56 ms, 1.65 s -> 1.34 s.

That one is not bit-identical - a multiply by a rounded reciprocal differs from a
correctly rounded quotient in the last place - so it can move a frame that sits
on the convergence tolerance. Battery: 21/24 space groups, no failures, and 17 of
24 crystals identical to the previous run, against a floor of 13 of 24 for the
same binary run twice.

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

image_analysis/indexing/FFTIndexerGPU.cu

@@ -91,6 +91,53 @@ __global__ void histogram_kernel(const float *__restrict__ coord_x,
     }
 }
 
+// The same histogram, one block per direction with the bins in shared memory.
+//
+// The kernel above gives a whole direction to a single thread, so neighbouring lanes write
+// histogram_size floats apart - 12.6 kB at the default sizing. Every warp instruction then touches 32
+// separate sectors of a buffer that is hundreds of megabytes (16384 directions x 3142 bins), with no
+// hope of staying in a 4 MB L2, and the projection loop becomes 160 million scattered global
+// read-modify-writes.
+//
+// The bins are counts, so here they are integers in shared memory. That matters twice: an integer
+// atomicAdd is a real shared-memory instruction on Turing and Ada, where the float one compiles to a
+// compare-and-swap retry loop; and a count below 2^24 converts to float exactly, so the output is bit
+// for bit what the repeated `+= 1.0` above produces. The division by histogram_spacing stays a
+// division - turning it into a multiply by the reciprocal would move a spot across a bin edge.
+__global__ void histogram_shared_kernel(const float *__restrict__ coord_x,
+                                        const float *__restrict__ coord_y,
+                                        const float *__restrict__ coord_z,
+                                        const float *__restrict__ dir_x,
+                                        const float *__restrict__ dir_y,
+                                        const float *__restrict__ dir_z,
+                                        float histogram_spacing,
+                                        int histogram_size,
+                                        int coord_size,
+                                        int direction_vectors_size,
+                                        float *__restrict__ output) {
+    extern __shared__ unsigned int bins[];
+    const int direction_idx = blockIdx.x;
+    if (direction_idx >= direction_vectors_size)
+        return;
+
+    for (int i = threadIdx.x; i < histogram_size; i += blockDim.x)
+        bins[i] = 0;
+    __syncthreads();
+
+    const float dx = dir_x[direction_idx], dy = dir_y[direction_idx], dz = dir_z[direction_idx];
+    for (int i = threadIdx.x; i < coord_size; i += blockDim.x) {
+        const float dot = fabsf(dx * coord_x[i] + dy * coord_y[i] + dz * coord_z[i]);
+        const int64_t bin = static_cast<int64_t>(dot / histogram_spacing);
+        if (bin >= 0 && bin < histogram_size)
+            atomicAdd(&bins[bin], 1u);
+    }
+    __syncthreads();
+
+    float *out = output + static_cast<size_t>(direction_idx) * static_cast<size_t>(histogram_size);
+    for (int i = threadIdx.x; i < histogram_size; i += blockDim.x)
+        out[i] = static_cast<float>(bins[i]);
+}
+
 inline void cuda_err(cudaError_t val) {
     if (val != cudaSuccess)
         throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
@@ -155,12 +202,21 @@ void FFTIndexerGPU::ExecuteFFT(const std::vector<Coord> &coord, size_t nspots) {
     cudaMemcpyAsync(d_spot_y, spot_y, nspots * sizeof(float), cudaMemcpyHostToDevice, stream);
     cudaMemcpyAsync(d_spot_z, spot_z, nspots * sizeof(float), cudaMemcpyHostToDevice, stream);
 
-    histogram_kernel<<<l_gridDim, l_blockDim, 0, stream>>>(d_spot_x, d_spot_y, d_spot_z,
-                                                           d_dir_x, d_dir_y, d_dir_z,
-                                                           histogram_spacing, histogram_size,
-                                                           nspots,
-                                                           direction_vectors.size(),
-                                                           d_input_fft);
+    // Shared-memory bins where they fit (they do at any sane sizing - 12.6 kB at the defaults), the
+    // thread-per-direction kernel where they do not.
+    const size_t hist_shared_bytes = static_cast<size_t>(histogram_size) * sizeof(unsigned int);
+    if (hist_shared_bytes <= 48 * 1024) {
+        histogram_shared_kernel<<<direction_vectors.size(), 256, hist_shared_bytes, stream>>>(
+                d_spot_x, d_spot_y, d_spot_z, d_dir_x, d_dir_y, d_dir_z,
+                histogram_spacing, histogram_size, nspots, direction_vectors.size(), d_input_fft);
+    } else {
+        histogram_kernel<<<l_gridDim, l_blockDim, 0, stream>>>(d_spot_x, d_spot_y, d_spot_z,
+                                                               d_dir_x, d_dir_y, d_dir_z,
+                                                               histogram_spacing, histogram_size,
+                                                               nspots,
+                                                               direction_vectors.size(),
+                                                               d_input_fft);
+    }
 
     cuda_err(cufftExecR2C(plan, d_input_fft, d_output_fft));
 

image_analysis/scale_merge/RotationScaleMerge.cpp

@@ -124,12 +124,12 @@ namespace {
             return 1.0;
         G = std::max(0.0, G);
 
-        const double k2 = robust_k * robust_k;
+        const double inv_k2 = 1.0 / (robust_k * robust_k);   // divisor hoisted; see the GPU kernel
         for (int iter = 0; iter < 30; ++iter) {
             const double G_prev = G;
             const double G_next = weighted_scale([&](const ScaleObs &o) {
                 const double res = o.weight * (G * o.coeff - o.Iobs);
-                return 1.0 / (1.0 + res * res / k2);
+                return 1.0 / (1.0 + res * res * inv_k2);
             });
             if (!std::isfinite(G_next))
                 break;

image_analysis/scale_merge/RotationScaleMergeGPU.cu

@@ -118,7 +118,10 @@ namespace {
         __syncthreads();
         if (s_cnt < MIN_REFLECTIONS) return;   // leave g[f]/scaled[f] as-is
 
-        const double k2 = robust_k * robust_k;
+        // 1/k^2, not k^2: the weight below divides by it once per observation per iteration, and the
+        // compiler will not hoist a loop-invariant divisor out of a double division - it emits the whole
+        // Newton refinement of the reciprocal every time. Same reason 1/sigma is precomputed above.
+        const double inv_k2 = 1.0 / (robust_k * robust_k);
 
         // seed: plain weighted-LS ratio (robust weight = 1)
         double num = 0.0, den = 0.0;
@@ -151,7 +154,7 @@ namespace {
                 const double w = inv_sigma[a];
                 const double w2 = w * w;
                 const double res = w * (G * coeff - double(I[a]));
-                const double rw = 1.0 / (1.0 + res * res / k2);
+                const double rw = 1.0 / (1.0 + res * res * inv_k2);
                 num += rw * w2 * coeff * double(I[a]);
                 den += rw * w2 * coeff * coeff;
             }

image_analysis/scale_merge/ScaleOnTheFly.cpp

@@ -57,12 +57,12 @@ namespace {
             return 1.0;
         G = std::max(0.0, G);
 
-        const double k2 = robust_k * robust_k;
+        const double inv_k2 = 1.0 / (robust_k * robust_k);   // divisor hoisted; see the GPU kernel
         for (int iter = 0; iter < 30; ++iter) {
             const double G_prev = G;
             const double G_next = weighted_scale([&](const ScaleObs &o) {
                 const double res = o.weight * (G * o.coeff - o.Iobs);
-                return 1.0 / (1.0 + res * res / k2);
+                return 1.0 / (1.0 + res * res * inv_k2);
             });
             if (!std::isfinite(G_next))
                 break;

image_analysis/scale_merge/StillsPartialityRefine.cpp

@@ -83,12 +83,12 @@ namespace {
             return 1.0;
         G = std::max(0.0, G);
 
-        const double k2 = robust_k * robust_k;
+        const double inv_k2 = 1.0 / (robust_k * robust_k);   // divisor hoisted; see the GPU kernel
         for (int iter = 0; iter < 30; ++iter) {
             const double G_prev = G;
             const double G_next = weighted_scale([&](size_t i) {
                 const double res = weight[i] * (G * coeff[i] - Iobs[i]);
-                return 1.0 / (1.0 + res * res / k2);
+                return 1.0 / (1.0 + res * res * inv_k2);
             });
             if (!std::isfinite(G_next))
                 break;