diff --git a/image_analysis/indexing/FFTIndexerGPU.cu b/image_analysis/indexing/FFTIndexerGPU.cu index 52ba31aa5..e6b011f2a 100644 --- a/image_analysis/indexing/FFTIndexerGPU.cu +++ b/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(dot / histogram_spacing); + if (bin >= 0 && bin < histogram_size) + atomicAdd(&bins[bin], 1u); + } + __syncthreads(); + + float *out = output + static_cast(direction_idx) * static_cast(histogram_size); + for (int i = threadIdx.x; i < histogram_size; i += blockDim.x) + out[i] = static_cast(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, 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<<>>(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(histogram_size) * sizeof(unsigned int); + if (hist_shared_bytes <= 48 * 1024) { + histogram_shared_kernel<<>>( + 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<<>>(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)); diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 8994d5d70..7cbc6aae8 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/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; diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 47cd4f83a..7effd9c5d 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/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; } diff --git a/image_analysis/scale_merge/ScaleOnTheFly.cpp b/image_analysis/scale_merge/ScaleOnTheFly.cpp index b1dc18872..132c812d7 100644 --- a/image_analysis/scale_merge/ScaleOnTheFly.cpp +++ b/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; diff --git a/image_analysis/scale_merge/StillsPartialityRefine.cpp b/image_analysis/scale_merge/StillsPartialityRefine.cpp index 6396d35eb..1d8aef081 100644 --- a/image_analysis/scale_merge/StillsPartialityRefine.cpp +++ b/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;