Give each FFT direction a block and its histogram shared memory

Cherry-picked from 2608-performance, restricted to the indexer: the same commit there also hoists a
reciprocal out of a loop in the scaling code, which is not being touched on this branch.

The histogram bins are unsigned integers and the counts stay below 2^24, so they convert to float
exactly - the vote is bit-identical, and the indexing result with it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VfYvJT5Nb71suJCowRBn5z

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));