diff --git a/image_analysis/indexing/FFTIndexerGPU.cu b/image_analysis/indexing/FFTIndexerGPU.cu index 52ba31aa..e6b011f2 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));