diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp index d0dbf5cb..a8163887 100644 --- a/image_analysis/IndexAndRefine.cpp +++ b/image_analysis/IndexAndRefine.cpp @@ -179,7 +179,10 @@ IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetry(Data int n = 0; for (const auto &q : recip) { const float hf = q * a, kf = q * b, lf = q * c; - const float dh = hf - std::round(hf), dk = kf - std::round(kf), dl = lf - std::round(lf); + // std::rint, not std::round: rounding half away from zero has to be a libm call, half to + // even is inlined. Only the squared residual is used, and the rules can differ only at an + // exact .5, where either leaves |frac| = 0.5 - so the count is the same either way. + const float dh = hf - std::rint(hf), dk = kf - std::rint(kf), dl = lf - std::rint(lf); if (dh * dh + dk * dk + dl * dl < idx_tol_sq) ++n; } const float frac = recip.empty() ? 0.0f : static_cast(n) / recip.size(); @@ -277,7 +280,9 @@ int CountIndexedSpots(const DiffractionGeometry &geom, const CrystalLattice &lat for (const auto &s : spots) { const Coord recip = s.ReciprocalCoord(geom); const float hf = recip * a, kf = recip * b, lf = recip * c; - const float dh = hf - std::round(hf), dk = kf - std::round(kf), dl = lf - std::round(lf); + // std::rint rather than std::round - only the squared residual is used, so the tie rule cannot + // change the count (see DetermineLatticeAndSymmetry). + const float dh = hf - std::rint(hf), dk = kf - std::rint(kf), dl = lf - std::rint(lf); if (dh * dh + dk * dk + dl * dl < tol_sq) ++n; } return n; diff --git a/image_analysis/MXAnalysisWithoutFPGA.cpp b/image_analysis/MXAnalysisWithoutFPGA.cpp index 5070dbf8..e4a54303 100644 --- a/image_analysis/MXAnalysisWithoutFPGA.cpp +++ b/image_analysis/MXAnalysisWithoutFPGA.cpp @@ -59,7 +59,11 @@ MXAnalysisWithoutFPGA::MXAnalysisWithoutFPGA(const DiffractionExperiment &in_exp } else { stream = std::make_shared(); preprocessor_buffer = std::make_unique(experiment.GetPixelsNum()); - preprocessor = std::make_unique(in_experiment, in_mask, stream); + // The preprocessed image only has to come back to the host if a CPU engine reads it. Every + // engine built below runs on the GPU, except the CPU adaptive finder that is kept when the fused + // GPU engine is off (the online receiver) - so that is the one case that needs the copy. + preprocessor = std::make_unique(in_experiment, in_mask, stream, + /*copy_image_to_host=*/!enable_fused_adaptive_gpu); spotFinder = std::make_unique(experiment.GetXPixelsNum(), experiment.GetYPixelsNum(), stream); azint = std::make_unique(integration, stream); bragg_engine = std::make_unique(in_experiment, stream); @@ -87,7 +91,7 @@ void MXAnalysisWithoutFPGA::Analyze(DataMessage &output, "Mismatch in pixel size"); const auto compression_start_time = std::chrono::steady_clock::now(); - const uint8_t *image_ptr = output.image.GetUncompressedPtr(decompression_buffer); + const uint8_t *image_ptr = Decompress(output.image); const auto compression_end_time = std::chrono::steady_clock::now(); if (output.image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION) output.compression_time_s = std::chrono::duration(compression_end_time - compression_start_time).count(); @@ -234,11 +238,19 @@ void MXAnalysisWithoutFPGA::AnalyzeROIOnly(DataMessage &output) { throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in pixel size"); - const uint8_t *image_ptr = output.image.GetUncompressedPtr(decompression_buffer); + const uint8_t *image_ptr = Decompress(output.image); preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode()); RunROIOnly(output); } +const uint8_t *MXAnalysisWithoutFPGA::Decompress(const CompressedImage &image) { + // An uncompressed image is read straight out of the message and never touches decompression_buffer, + // so it stays in pageable memory - the buffer is only worth page-locking when it is actually used. + if (image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION) + preprocessor->PinInputBuffer(decompression_buffer, image.GetUncompressedSize()); + return image.GetUncompressedPtr(decompression_buffer); +} + void MXAnalysisWithoutFPGA::RunROIOnly(DataMessage &output) { output.roi.clear(); if (roi) diff --git a/image_analysis/MXAnalysisWithoutFPGA.h b/image_analysis/MXAnalysisWithoutFPGA.h index 64eb08f7..a018c0b1 100644 --- a/image_analysis/MXAnalysisWithoutFPGA.h +++ b/image_analysis/MXAnalysisWithoutFPGA.h @@ -53,6 +53,10 @@ class MXAnalysisWithoutFPGA { std::unique_ptr preprocessor_buffer; const PixelMask &mask; + // Decompress the image into decompression_buffer (or read it straight from the message, when it is + // not compressed) and return where it landed. + const uint8_t *Decompress(const CompressedImage &image); + std::vector mask_resolution; // The limits mask_resolution was built for. Kept as the OPTIONAL the caller passed, so an unset // high-resolution limit compares equal to itself and the mask is not rebuilt on every image. diff --git a/image_analysis/geom_refinement/XtalOptimizer.cpp b/image_analysis/geom_refinement/XtalOptimizer.cpp index a2843d78..348cc9b0 100644 --- a/image_analysis/geom_refinement/XtalOptimizer.cpp +++ b/image_analysis/geom_refinement/XtalOptimizer.cpp @@ -185,6 +185,17 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, const float tolerance_sq = tolerance * tolerance; + // Per-image rotation refinement frees only the beam and the orientation and holds the other five + // blocks constant, so the seven-block residual makes Ceres differentiate 17 parameters to use 5. + // Where that is the configuration, use the reduced residual instead - identical fit, Jet<5> + // autodiff. Any other combination (stills also free the cell, the offline refiner frees distance + // and detector angles) keeps the general form below. + const bool beam_and_orientation_only = data.refine_beam_center + && !data.refine_distance_mm + && !data.refine_detector_angles + && !data.refine_rotation_axis + && !data.refine_unit_cell; + // Sum of w^2 over the spots that entered - the beam prior below is scaled by it so that its // strength relative to the data is the same weighted or not. Equals the residual block count // when the spots are unweighted. @@ -234,42 +245,49 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, const double weight_sq = weight.empty() ? 1.0 : weight[j] * weight[j]; effective_spots += weight_sq; - problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( - new XtalResidual(pt.x, pt.y, - data.geom.GetWavelength_A(), - data.geom.GetPixelSize_mm(), - data.geom.GetPoniRot3_rad(), - angle_rad, - h, k, l, - data.crystal_system)), - // Ceres has no per-residual weight; ScaledLoss(nullptr, a) multiplies the squared - // residual by the constant a, i.e. it applies a weight of sqrt(a) to the residual. - weight.empty() - ? nullptr - : new ceres::ScaledLoss(nullptr, weight_sq, ceres::TAKE_OWNERSHIP), - beam, - &distance_mm, - detector_rot, - rot_vec, - latt_vec0, - latt_vec1, - latt_vec2 - ); + const XtalResidual residual(pt.x, pt.y, + data.geom.GetWavelength_A(), + data.geom.GetPixelSize_mm(), + data.geom.GetPoniRot3_rad(), + angle_rad, + h, k, l, + data.crystal_system); + + // Ceres has no per-residual weight; ScaledLoss(nullptr, a) multiplies the squared + // residual by the constant a, i.e. it applies a weight of sqrt(a) to the residual. + ceres::LossFunction *loss = weight.empty() + ? nullptr + : new ceres::ScaledLoss(nullptr, weight_sq, + ceres::TAKE_OWNERSHIP); + + if (beam_and_orientation_only) + problem.AddResidualBlock( + new ceres::AutoDiffCostFunction( + new XtalResidualBeamOrientation(residual, distance_mm, detector_rot, + rot_vec, latt_vec1, latt_vec2)), + loss, + beam, + latt_vec0 + ); + else + problem.AddResidualBlock( + new ceres::AutoDiffCostFunction( + new XtalResidual(residual)), + loss, + beam, + &distance_mm, + detector_rot, + rot_vec, + latt_vec0, + latt_vec1, + latt_vec2 + ); } } if (problem.NumResidualBlocks() < data.min_spots) return false; - if (!data.refine_distance_mm) - problem.SetParameterBlockConstant(&distance_mm); - else { - const double dist_range = 0.1; - problem.SetParameterLowerBound(&distance_mm, 0, distance_mm * (1.0 - dist_range)); - problem.SetParameterUpperBound(&distance_mm, 0, distance_mm * (1.0 + dist_range)); - } - if (!data.refine_beam_center) problem.SetParameterBlockConstant(beam); else if (data.axis) { @@ -296,48 +314,60 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, nullptr, beam); } - if (!data.refine_detector_angles) { - problem.SetParameterBlockConstant(detector_rot); - } else { - const double rot_range = 3.0 / 180.0 * PI; - for (int i = 0; i < 2; ++i) { - problem.SetParameterLowerBound(detector_rot, i, detector_rot[i] - rot_range); - problem.SetParameterUpperBound(detector_rot, i, detector_rot[i] + rot_range); - } - } - - if (!data.refine_rotation_axis) { - problem.SetParameterBlockConstant(rot_vec); - } - - if (!data.refine_unit_cell) { - problem.SetParameterBlockConstant(latt_vec1); - problem.SetParameterBlockConstant(latt_vec2); - } else { - // Parameter bounds - // Lengths - for (int i = 0; i < 3; ++i) { - problem.SetParameterLowerBound(latt_vec1, i, data.min_length_A); - problem.SetParameterUpperBound(latt_vec1, i, data.max_length_A); + // Distance, detector angles, rotation axis and cell are parameter blocks only in the general + // seven-block residual; the reduced one bakes them in, so there is nothing left to configure. + if (!beam_and_orientation_only) { + if (!data.refine_distance_mm) + problem.SetParameterBlockConstant(&distance_mm); + else { + const double dist_range = 0.1; + problem.SetParameterLowerBound(&distance_mm, 0, distance_mm * (1.0 - dist_range)); + problem.SetParameterUpperBound(&distance_mm, 0, distance_mm * (1.0 + dist_range)); } - if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) { - const double beta_lo = std::max(1e-6, PI * (data.min_angle_deg / 180.0)); - const double beta_hi = std::min(PI - 1e-6, PI * (data.max_angle_deg / 180.0)); - problem.SetParameterLowerBound(latt_vec2, 0, beta_lo); - problem.SetParameterUpperBound(latt_vec2, 0, beta_hi); - } else if (data.crystal_system == gemmi::CrystalSystem::Triclinic) { - // α, β, γ bounds (radians) - const double alo = PI * (data.min_angle_deg / 180.0); - const double ahi = PI * (data.max_angle_deg / 180.0); - for (int i = 0; i < 3; ++i) { - problem.SetParameterLowerBound(latt_vec2, i, alo); - problem.SetParameterUpperBound(latt_vec2, i, ahi); - } + if (!data.refine_detector_angles) { + problem.SetParameterBlockConstant(detector_rot); } else { - // Orthorhombic / Tetragonal / Cubic / Hexagonal: - // latt_vec2 has no meaning for these systems — always freeze it. + const double rot_range = 3.0 / 180.0 * PI; + for (int i = 0; i < 2; ++i) { + problem.SetParameterLowerBound(detector_rot, i, detector_rot[i] - rot_range); + problem.SetParameterUpperBound(detector_rot, i, detector_rot[i] + rot_range); + } + } + + if (!data.refine_rotation_axis) { + problem.SetParameterBlockConstant(rot_vec); + } + + if (!data.refine_unit_cell) { + problem.SetParameterBlockConstant(latt_vec1); problem.SetParameterBlockConstant(latt_vec2); + } else { + // Parameter bounds + // Lengths + for (int i = 0; i < 3; ++i) { + problem.SetParameterLowerBound(latt_vec1, i, data.min_length_A); + problem.SetParameterUpperBound(latt_vec1, i, data.max_length_A); + } + + if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) { + const double beta_lo = std::max(1e-6, PI * (data.min_angle_deg / 180.0)); + const double beta_hi = std::min(PI - 1e-6, PI * (data.max_angle_deg / 180.0)); + problem.SetParameterLowerBound(latt_vec2, 0, beta_lo); + problem.SetParameterUpperBound(latt_vec2, 0, beta_hi); + } else if (data.crystal_system == gemmi::CrystalSystem::Triclinic) { + // α, β, γ bounds (radians) + const double alo = PI * (data.min_angle_deg / 180.0); + const double ahi = PI * (data.max_angle_deg / 180.0); + for (int i = 0; i < 3; ++i) { + problem.SetParameterLowerBound(latt_vec2, i, alo); + problem.SetParameterUpperBound(latt_vec2, i, ahi); + } + } else { + // Orthorhombic / Tetragonal / Cubic / Hexagonal: + // latt_vec2 has no meaning for these systems — always freeze it. + problem.SetParameterBlockConstant(latt_vec2); + } } } diff --git a/image_analysis/geom_refinement/XtalResidual.h b/image_analysis/geom_refinement/XtalResidual.h index 6bdcbc3c..cf6a5a4d 100644 --- a/image_analysis/geom_refinement/XtalResidual.h +++ b/image_analysis/geom_refinement/XtalResidual.h @@ -212,3 +212,42 @@ struct XtalResidual { const double angle_rad; gemmi::CrystalSystem symmetry; }; + +// Same residual with distance, detector angles, rotation axis and unit cell baked in as constants, so +// that only beam(2) and orientation(3) remain parameter blocks. Ceres sizes its autodiff dual numbers +// from the DECLARED blocks, not from which of them the caller then holds constant, so the seven-block +// form above differentiates 17 parameters even when 5 are free; this one runs on Jet. It +// forwards to XtalResidual with the very same values, so the residual and the beam/orientation columns +// of its Jacobian are unchanged. +struct XtalResidualBeamOrientation { + XtalResidualBeamOrientation(const XtalResidual &residual, + double distance_mm, + const double *detector_rot, + const double *rotation_axis, + const double *uc_len, + const double *uc_angle) + : residual(residual), + distance_mm(distance_mm), + detector_rot{detector_rot[0], detector_rot[1]}, + rotation_axis{rotation_axis[0], rotation_axis[1], rotation_axis[2]}, + uc_len{uc_len[0], uc_len[1], uc_len[2]}, + uc_angle{uc_angle[0], uc_angle[1], uc_angle[2]} { + } + + template + bool operator()(const T *const beam, const T *const p0, T *residual_out) const { + const T distance[1] = {T(distance_mm)}; + const T rot[2] = {T(detector_rot[0]), T(detector_rot[1])}; + const T axis[3] = {T(rotation_axis[0]), T(rotation_axis[1]), T(rotation_axis[2])}; + const T p1[3] = {T(uc_len[0]), T(uc_len[1]), T(uc_len[2])}; + const T p2[3] = {T(uc_angle[0]), T(uc_angle[1]), T(uc_angle[2])}; + return residual(beam, distance, rot, axis, p0, p1, p2, residual_out); + } + + const XtalResidual residual; + const double distance_mm; + const double detector_rot[2]; + const double rotation_axis[3]; + const double uc_len[3]; + const double uc_angle[3]; +}; diff --git a/image_analysis/image_preprocessing/ImagePreprocessor.h b/image_analysis/image_preprocessing/ImagePreprocessor.h index e105bf86..3dd8af47 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessor.h +++ b/image_analysis/image_preprocessing/ImagePreprocessor.h @@ -5,6 +5,7 @@ #include #include +#include #include "../common/CompressedImage.h" #include "../common/DiffractionExperiment.h" #include "../common/PixelMask.h" @@ -27,4 +28,10 @@ public: ImagePreprocessor(const DiffractionExperiment &experiment); virtual ~ImagePreprocessor() = default; virtual ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *decompressed_image, CompressedImageMode image_mode) = 0; + + // Resize the buffer an image will be decompressed into and page-lock it, so that the host->device + // copy of Analyze() is a real DMA. Without page-locking the driver stages the copy through its own + // pinned pool, which is a host-side copy on the calling thread: it does not overlap and it degrades + // badly with the number of workers. Nothing to do on the CPU. + virtual void PinInputBuffer(std::vector &buffer, size_t size) {} }; diff --git a/image_analysis/image_preprocessing/ImagePreprocessorBuffer.cpp b/image_analysis/image_preprocessing/ImagePreprocessorBuffer.cpp index 421259be..abdb2c53 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorBuffer.cpp +++ b/image_analysis/image_preprocessing/ImagePreprocessorBuffer.cpp @@ -34,6 +34,12 @@ const int32_t *ImagePreprocessorBuffer::data() const { return buffer.data(); } +void ImagePreprocessorBuffer::Gather(const std::vector &npixel, std::vector &values) const { + values.resize(npixel.size()); + for (size_t i = 0; i < npixel.size(); i++) + values[i] = buffer[npixel[i]]; +} + int32_t *ImagePreprocessorBuffer::getGPUBuffer() { return nullptr; } diff --git a/image_analysis/image_preprocessing/ImagePreprocessorBuffer.h b/image_analysis/image_preprocessing/ImagePreprocessorBuffer.h index f380ca1d..3644ff03 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorBuffer.h +++ b/image_analysis/image_preprocessing/ImagePreprocessorBuffer.h @@ -27,4 +27,8 @@ public: // GPU operations (overriden in ImagePreprocessorBufferGPU virtual int32_t *getGPUBuffer(); virtual const int32_t *getGPUBuffer() const; + + // Values of the pixels with the given flat indices. Not just operator[] in a loop: on the GPU the + // preprocessed image can live in device memory only, and then this gathers them there. + virtual void Gather(const std::vector &npixel, std::vector &values) const; }; diff --git a/image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.cu b/image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.cu index 465e4be3..ed02bd30 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.cu +++ b/image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.cu @@ -3,10 +3,20 @@ #include "ImagePreprocessorBufferGPU.h" +__global__ void gather_kernel(const int32_t *__restrict__ image, + const uint32_t *__restrict__ npixel, + int32_t *__restrict__ values, + int count) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < count; i += blockDim.x * gridDim.x) + values[i] = image[npixel[i]]; +} + ImagePreprocessorBufferGPU::ImagePreprocessorBufferGPU(size_t npixel) : ImagePreprocessorBuffer(npixel), gpu_image(npixel), - buffer_reg(buffer) { + buffer_reg(buffer), + gpu_gather_index(MAX_GATHER), + gpu_gather_value(MAX_GATHER) { } int32_t *ImagePreprocessorBufferGPU::getGPUBuffer() { @@ -16,3 +26,18 @@ int32_t *ImagePreprocessorBufferGPU::getGPUBuffer() { const int32_t *ImagePreprocessorBufferGPU::getGPUBuffer() const { return gpu_image; } + +void ImagePreprocessorBufferGPU::Gather(const std::vector &npixel, std::vector &values) const { + values.resize(npixel.size()); + if (npixel.empty()) + return; + + const int count = static_cast(npixel.size()); + cudaMemcpyAsync(gpu_gather_index.get(), npixel.data(), count * sizeof(uint32_t), + cudaMemcpyHostToDevice, gather_stream); + gather_kernel<<<(count + 255) / 256, 256, 0, gather_stream>>>( + gpu_image.get(), gpu_gather_index.get(), gpu_gather_value.get(), count); + cudaMemcpyAsync(values.data(), gpu_gather_value.get(), count * sizeof(int32_t), + cudaMemcpyDeviceToHost, gather_stream); + cudaStreamSynchronize(gather_stream); +} diff --git a/image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h b/image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h index 6db2c780..d03a94aa 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h +++ b/image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h @@ -12,9 +12,20 @@ class ImagePreprocessorBufferGPU : public ImagePreprocessorBuffer { CudaDevicePtr gpu_image; CudaRegisteredVector buffer_reg; + // Staging for Gather(). Its only caller is ImageSpotFinder::ExtractSpots, which gives up on a frame + // with UINT16_MAX or more strong pixels (the connected-component search rejects it anyway), so that + // is the largest gather that can be asked for. + static constexpr size_t MAX_GATHER = UINT16_MAX; + CudaDevicePtr gpu_gather_index; + CudaDevicePtr gpu_gather_value; + // Own stream: every analysis engine synchronises its own stream before it returns, so the device + // image is final by the time a gather is asked for. The NULL stream would serialise all workers. + CudaStream gather_stream; + public: explicit ImagePreprocessorBufferGPU(size_t npixel); int32_t *getGPUBuffer() override; const int32_t *getGPUBuffer() const override; + void Gather(const std::vector &npixel, std::vector &values) const override; }; diff --git a/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu b/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu index 8844134d..b8f74b84 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu +++ b/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu @@ -90,9 +90,10 @@ __global__ void preprocess_kernel( } ImagePreprocessorGPU::ImagePreprocessorGPU(const DiffractionExperiment &experiment, const PixelMask &mask, - std::shared_ptr stream) + std::shared_ptr stream, bool copy_image_to_host) : ImagePreprocessor(experiment), stream(stream), + copy_image_to_host(copy_image_to_host), gpu_decompressed_image(npixels * sizeof(uint32_t)), // Overshoot - if input image is 1- or 2-byte, then it is still fine, while memory loss is minimal gpu_stats(1), cpu_stats(1), @@ -115,6 +116,15 @@ ImagePreprocessorGPU::ImagePreprocessorGPU(const DiffractionExperiment &experime blocks = 4 * prop.multiProcessorCount; } +void ImagePreprocessorGPU::PinInputBuffer(std::vector &buffer, size_t size) { + if (buffer.size() == size) + return; + // Unregister before the resize, which can move the buffer. + input_reg.unregister(); + buffer.resize(size); + input_reg.rebind(buffer); +} + ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *image_ptr, CompressedImageMode image_mode) { switch (image_mode) { @@ -158,7 +168,10 @@ ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed sat_value, err_value, npixels); - cudaMemcpyAsync(processed_image.data(), processed_image.getGPUBuffer(), npixels * sizeof(int32_t), cudaMemcpyDeviceToHost, *stream); + // 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. + if (copy_image_to_host) + cudaMemcpyAsync(processed_image.data(), processed_image.getGPUBuffer(), npixels * sizeof(int32_t), cudaMemcpyDeviceToHost, *stream); cudaMemcpyAsync(cpu_stats.data(), gpu_stats, sizeof(ImageStatistics), cudaMemcpyDeviceToHost, *stream); cudaStreamSynchronize(*stream); diff --git a/image_analysis/image_preprocessing/ImagePreprocessorGPU.h b/image_analysis/image_preprocessing/ImagePreprocessorGPU.h index 3d4d3b92..9dbfb646 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorGPU.h +++ b/image_analysis/image_preprocessing/ImagePreprocessorGPU.h @@ -9,6 +9,7 @@ class ImagePreprocessorGPU : public ImagePreprocessor { std::shared_ptr stream; + const bool copy_image_to_host; int threads; int blocks; // Geometry-only, so one copy per GPU shared with every other engine on it (CudaSharedTables.h). @@ -18,12 +19,18 @@ class ImagePreprocessorGPU : public ImagePreprocessor { std::vector cpu_stats; CudaRegisteredVector cpu_stats_reg; + CudaRegisteredVector input_reg; // page-locks the caller's decompression buffer std::vector cpu_image; template ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *input, T err_value, T sat_value); public: - ImagePreprocessorGPU(const DiffractionExperiment &experiment, const PixelMask &mask, std::shared_ptr stream); + // copy_image_to_host copies the preprocessed image back after every frame. It is only needed when + // something on the CPU reads it - the GPU engines all work off the device buffer - and at 4 bytes + // per pixel it is the single largest transfer in the pipeline, so the caller says whether it wants it. + ImagePreprocessorGPU(const DiffractionExperiment &experiment, const PixelMask &mask, std::shared_ptr stream, + bool copy_image_to_host = true); ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *decompressed_image, CompressedImageMode image_mode) override; + void PinInputBuffer(std::vector &buffer, size_t size) override; }; diff --git a/image_analysis/indexing/AnalyzeIndexing.cpp b/image_analysis/indexing/AnalyzeIndexing.cpp index af17244a..6c4eeeea 100644 --- a/image_analysis/indexing/AnalyzeIndexing.cpp +++ b/image_analysis/indexing/AnalyzeIndexing.cpp @@ -340,25 +340,32 @@ bool AnalyzeIndexing(DataMessage &message, float k_fp = recip * b; float l_fp = recip * c; - float h_frac = h_fp - std::round(h_fp); - float k_frac = k_fp - std::round(k_fp); - float l_frac = l_fp - std::round(l_fp); + // std::rint, not std::round: rounding half away from zero has to be a libm call, rounding half + // to even is a handful of inline instructions. Only the SQUARED residual is taken here, and the + // two rules can differ only at an exact .5, where either leaves |frac| = 0.5 - so norm_sq is the + // same number. The Miller index itself keeps std::round, below, and is only paid for when the + // spot actually indexes. + float h_frac = h_fp - std::rint(h_fp); + float k_frac = k_fp - std::rint(k_fp); + float l_frac = l_fp - std::rint(l_fp); float norm_sq = h_frac * h_frac + k_frac * k_frac + l_frac * l_frac; - Coord recip_pred = std::round(h_fp) * astar + std::round(k_fp) * bstar + std::round(l_fp) * cstar; - // See indexing_peak_check() in peaks.c in CrystFEL if (norm_sq < indexing_tolerance_sq) { if (index_ice_ring || !message.spots[i].ice_ring) { nspots_indexed++; indexed_q_sq_max = std::max(indexed_q_sq_max, spot_q_sq[i]); } + const float h_r = std::round(h_fp); + const float k_r = std::round(k_fp); + const float l_r = std::round(l_fp); + Coord recip_pred = h_r * astar + k_r * bstar + l_r * cstar; indexed_spots[i] = 1; message.spots[i].dist_ewald_sphere = geom.DistFromEwaldSphere(recip_pred); - message.spots[i].h = std::lround(h_fp); - message.spots[i].k = std::lround(k_fp); - message.spots[i].l = std::lround(l_fp); + message.spots[i].h = static_cast(h_r); + message.spots[i].k = static_cast(k_r); + message.spots[i].l = static_cast(l_r); } } @@ -428,22 +435,24 @@ bool AnalyzeIndexing(DataMessage &message, float k_fp = recip * eb; float l_fp = recip * ec; - float h_frac = h_fp - std::round(h_fp); - float k_frac = k_fp - std::round(k_fp); - float l_frac = l_fp - std::round(l_fp); + // std::rint for the residual, std::round for the index - see the main-lattice loop. + float h_frac = h_fp - std::rint(h_fp); + float k_frac = k_fp - std::rint(k_fp); + float l_frac = l_fp - std::rint(l_fp); float norm_sq = h_frac * h_frac + k_frac * k_frac + l_frac * l_frac; if (norm_sq < indexing_tolerance_sq) { - Coord recip_pred = std::round(h_fp) * east - + std::round(k_fp) * ebst - + std::round(l_fp) * ecst; + const float h_r = std::round(h_fp); + const float k_r = std::round(k_fp); + const float l_r = std::round(l_fp); + Coord recip_pred = h_r * east + k_r * ebst + l_r * ecst; message.spots[i].indexed = true; message.spots[i].lattice = lattice_id; message.spots[i].dist_ewald_sphere = geom.DistFromEwaldSphere(recip_pred); - message.spots[i].h = std::lround(h_fp); - message.spots[i].k = std::lround(k_fp); - message.spots[i].l = std::lround(l_fp); + message.spots[i].h = static_cast(h_r); + message.spots[i].k = static_cast(k_r); + message.spots[i].l = static_cast(l_r); } } diff --git a/image_analysis/indexing/FFTIndexer.cpp b/image_analysis/indexing/FFTIndexer.cpp index 86cabc4e..273ff977 100644 --- a/image_analysis/indexing/FFTIndexer.cpp +++ b/image_analysis/indexing/FFTIndexer.cpp @@ -281,9 +281,13 @@ float FFTIndexer::IndexedFraction(const CrystalLattice &latt, size_t indexed = 0; for (size_t i = 0; i < nspots; i++) { const Coord &s = coord[i]; - const float dh = a * s - std::round(a * s); // Coord operator* = dot product = Miller index - const float dk = b * s - std::round(b * s); - const float dl = c * s - std::round(c * s); + const float hf = a * s, kf = b * s, lf = c * s; // Coord operator* = dot product = Miller index + // std::rint, not std::round: rounding half away from zero has to be a libm call, half to even is + // inlined. Only the squared residual is used, and the rules can differ only at an exact .5, where + // either leaves |frac| = 0.5 - so the indexed count is the same either way. + const float dh = hf - std::rint(hf); + const float dk = kf - std::rint(kf); + const float dl = lf - std::rint(lf); if (dh * dh + dk * dk + dl * dl < tol_sq) ++indexed; } diff --git a/image_analysis/indexing/PostIndexingRefinement.cpp b/image_analysis/indexing/PostIndexingRefinement.cpp index 029241f6..8725c75f 100644 --- a/image_analysis/indexing/PostIndexingRefinement.cpp +++ b/image_analysis/indexing/PostIndexingRefinement.cpp @@ -44,9 +44,14 @@ namespace { int64_t &indexed_spot_count) { const float indexing_tolerance_sq = indexing_tolerance * indexing_tolerance; - // Compute fractional Miller indices + // Compute fractional Miller indices. rint (round half to even) rather than round (round half + // away from zero): without SSE4.1 Eigen has no vector round, so each element is a libm call, + // while rint is a few inline instructions. Only the SQUARED residual is taken below and the two + // rules can differ only at an exact .5, where either leaves |frac| = 0.5 - so the mask and the + // count are the same. The refinement loop above keeps round: there the rounded value IS the + // Miller index that goes into the residual and the QR solve, so its tie rule does matter. Eigen::MatrixX3 miller_frac = spots * cell; - Eigen::MatrixX3 miller_int = miller_frac.array().round().matrix(); + Eigen::MatrixX3 miller_int = miller_frac.array().rint().matrix(); Eigen::MatrixX3 frac_resid = miller_frac - miller_int; std::vector mask(spots.rows(), 0); diff --git a/image_analysis/scale_merge/FrenchWilson.cpp b/image_analysis/scale_merge/FrenchWilson.cpp index bbb7c169..362e42b3 100644 --- a/image_analysis/scale_merge/FrenchWilson.cpp +++ b/image_analysis/scale_merge/FrenchWilson.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include @@ -22,12 +23,13 @@ struct Posterior { // prior with mean sigma_wilson. Integrated numerically over J in [0, I + 8 sigma] with a log-shift // so the exponentials never overflow/underflow. acentric: p(J) ~ exp(-J/S); centric: // p(J) ~ exp(-J/2S)/sqrt(J). -Posterior integrate_posterior(double I, double sigma, double sigma_wilson, bool centric, int npts) { +// `logw` is caller-owned scratch of npts doubles (one per worker), so the integration allocates nothing. +Posterior integrate_posterior(double I, double sigma, double sigma_wilson, bool centric, int npts, + std::vector &logw) { const double inv_2s2 = 1.0 / (2.0 * sigma * sigma); const double j_max = std::max(I, 0.0) + 8.0 * sigma; const double dj = j_max / npts; - std::vector logw(npts); double max_logw = -std::numeric_limits::infinity(); for (int i = 0; i < npts; ++i) { const double j = (i + 0.5) * dj; @@ -120,20 +122,43 @@ void ApplyFrenchWilson(std::vector &merged, int32_t space_grou // French-Wilson |F| for one intensity of reflection r (its mean, or one Bijvoet hand); the shell // Wilson prior, epsilon and centric flag are the reflection's, shared by all three. - auto fw_one = [&](const MergedReflection &r, float I, float sigma, float &F, float &sigF) { + auto fw_one = [&](const MergedReflection &r, float I, float sigma, float &F, float &sigF, + std::vector &logw) { if (!std::isfinite(I) || !std::isfinite(sigma) || sigma <= 0.0f) { naive_one(I, sigma, F, sigF); return; } // Strong reflections: the FW correction is negligible, <|F|> = sqrt(I). if (I > opts.strong_cutoff * sigma) { naive_one(I, sigma, F, sigF); return; } const auto s = shells.GetShell(r.d); const double sigma_wilson = epsilon(r) * (s ? shell_mean[*s] : global_mean); const bool centric = gops.is_reflection_centric({{r.h, r.k, r.l}}); - const Posterior post = integrate_posterior(I, sigma, sigma_wilson, centric, opts.integration_points); + const Posterior post = integrate_posterior(I, sigma, sigma_wilson, centric, + opts.integration_points, logw); F = static_cast(post.mean_F); sigF = static_cast(std::sqrt(std::max(0.0, post.mean_I - post.mean_F * post.mean_F))); }; - for (auto &r : merged) { - fw_one(r, r.I, r.sigma, r.F, r.sigmaF); - fw_one(r, r.I_plus, r.sigma_plus, r.F_plus, r.sigmaF_plus); - fw_one(r, r.I_minus, r.sigma_minus, r.F_minus, r.sigmaF_minus); + // Each reflection's amplitudes depend only on itself and the shell priors above, so the loop is + // data-parallel over contiguous chunks and gives the same result whatever the worker count. + const int n = static_cast(merged.size()); + const int nt = std::clamp(opts.num_threads, 1, n); + const int chunk = (n + nt - 1) / nt; + auto do_chunk = [&](int lo, int hi) { + std::vector logw(opts.integration_points); + for (int i = lo; i < hi; ++i) { + MergedReflection &r = merged[i]; + fw_one(r, r.I, r.sigma, r.F, r.sigmaF, logw); + fw_one(r, r.I_plus, r.sigma_plus, r.F_plus, r.sigmaF_plus, logw); + fw_one(r, r.I_minus, r.sigma_minus, r.F_minus, r.sigmaF_minus, logw); + } + }; + if (nt == 1) { + do_chunk(0, n); + return; } + std::vector> futures; + futures.reserve(nt); + for (int t = 0; t < nt; ++t) { + const int lo = t * chunk, hi = std::min(n, lo + chunk); + if (lo >= hi) break; + futures.emplace_back(std::async(std::launch::async, [&do_chunk, lo, hi] { do_chunk(lo, hi); })); + } + for (auto &f : futures) f.get(); } diff --git a/image_analysis/scale_merge/FrenchWilson.h b/image_analysis/scale_merge/FrenchWilson.h index 9ea38350..fd41e020 100644 --- a/image_analysis/scale_merge/FrenchWilson.h +++ b/image_analysis/scale_merge/FrenchWilson.h @@ -13,6 +13,7 @@ struct FrenchWilsonOptions { int min_reflections_per_shell = 20; // below this, fall back to the overall mean for the shell int integration_points = 400; // posterior integration steps (weak reflections only) double strong_cutoff = 4.0; // I/sigma above which <|F|> = sqrt(I) (FW bias negligible) + int num_threads = 1; // workers for the per-reflection integration (independent per reflection) }; // Fill F and sigmaF on every merged reflection with the French-Wilson estimate of the amplitude: diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 37d8313d..47208ed6 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -248,33 +248,45 @@ void RotationScaleMerge::Ingest() { // Sort ONCE by (raw h,k,l, image_number) and split into raw-hkl runs. This is the one expensive sort; // both the 3D combine (event split) and the per-space-group ASU grouping reuse this order. + // Sorting an index array whose comparator dereferences the 72-byte Obs is a cache miss per + // comparison over a multi-GB array, so sort a packed copy of the key fields instead. Started from + // the identity order with the same comparisons in the same order, std::sort takes exactly the same + // branches and produces exactly the same permutation. + struct SortKey { int32_t h, k, l; float image_number; int32_t idx; }; perm.resize(partials.size()); - for (int i = 0; i < static_cast(partials.size()); ++i) perm[i] = i; - std::sort(perm.begin(), perm.end(), [&](int32_t a, int32_t b) { - const auto &x1 = partials[a]; - const auto &y1 = partials[b]; - if (x1.h != y1.h) return x1.h < y1.h; - if (x1.k != y1.k) return x1.k < y1.k; - if (x1.l != y1.l) return x1.l < y1.l; - return x1.image_number < y1.image_number; - }); rawrun_start.clear(); rawrun_count.clear(); rawrun_h.clear(); rawrun_k.clear(); rawrun_l.clear(); rawrun_d.clear(); - for (int i = 0; i < static_cast(perm.size()); ) { - const auto &o0 = partials[perm[i]]; - int j = i; - float d = NAN; - while (j < static_cast(perm.size())) { - const auto &o = partials[perm[j]]; - if (o.h != o0.h || o.k != o0.k || o.l != o0.l) break; - if (!std::isfinite(d) && std::isfinite(o.d) && o.d > 0.0f) d = o.d; - ++j; + { + std::vector keys(partials.size()); + for (size_t i = 0; i < partials.size(); ++i) { + const auto &o = partials[i]; + keys[i] = SortKey{o.h, o.k, o.l, o.image_number, static_cast(i)}; + } + std::sort(keys.begin(), keys.end(), [](const SortKey &a, const SortKey &b) { + if (a.h != b.h) return a.h < b.h; + if (a.k != b.k) return a.k < b.k; + if (a.l != b.l) return a.l < b.l; + return a.image_number < b.image_number; + }); + for (size_t i = 0; i < keys.size(); ++i) perm[i] = keys[i].idx; + for (int i = 0; i < static_cast(keys.size()); ) { + const SortKey k0 = keys[i]; + int j = i; + float d = NAN; + while (j < static_cast(keys.size())) { + if (keys[j].h != k0.h || keys[j].k != k0.k || keys[j].l != k0.l) break; + if (!std::isfinite(d)) { + const float dj = partials[keys[j].idx].d; // resolution: only until one is usable + if (std::isfinite(dj) && dj > 0.0f) d = dj; + } + ++j; + } + rawrun_start.push_back(i); + rawrun_count.push_back(j - i); + rawrun_h.push_back(k0.h); rawrun_k.push_back(k0.k); rawrun_l.push_back(k0.l); + rawrun_d.push_back(d); + i = j; } - rawrun_start.push_back(i); - rawrun_count.push_back(j - i); - rawrun_h.push_back(o0.h); rawrun_k.push_back(o0.k); rawrun_l.push_back(o0.l); - rawrun_d.push_back(d); - i = j; } rawrun_group.assign(rawrun_start.size(), -1); logger.Info("RotationScaleMerge: ingested {} partial observations from {} frames ({} distinct hkl)", @@ -960,24 +972,42 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int auto usable = [&](const Obs &o) { return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && o.partiality >= min_partiality; }; + // The observations any of the passes below can use, in fulls order, split by frame parity. Every + // pass then walks one list instead of re-testing (and re-reading) the whole of fulls - the + // cross-validated halves cost half a pass rather than a whole one with a skip. The lists keep the + // fulls order, so each sum below is formed from exactly the same terms in exactly the same order. + std::vector idx_all, idx_even, idx_odd; + idx_all.reserve(fulls.size()); + for (size_t i = 0; i < fulls.size(); ++i) { + if (!usable(fulls[i]) || cell[i] < 0) continue; + idx_all.push_back(static_cast(i)); + (fulls[i].frame & 1 ? idx_odd : idx_even).push_back(static_cast(i)); + } + auto subset = [&](int parity) -> const std::vector & { + return parity < 0 ? idx_all : (parity ? idx_odd : idx_even); + }; + // Per-group reference accumulators, reused by every pass (n_groups is large, so re-allocating them + // ~22x per surface is ~22x the page faults for nothing). + std::vector sw(n_groups), swI(n_groups); + // Fit the per-cell factor over the subset {frame&1 == parity} (parity < 0 = all fulls), n_iter // alternating rounds against that subset's own reference (Tikhonov pull to 1, gauge-fixed to a // den-weighted geometric mean of 1 so it never drifts the overall scale). Returns the per-cell factor. auto fit_surface = [&](int parity) -> std::vector { + const std::vector &sel = subset(parity); std::vector A(ncell, 1.0); for (int it = 0; it < n_iter; ++it) { - std::vector sw(n_groups, 0.0), swI(n_groups, 0.0); - for (size_t i = 0; i < fulls.size(); ++i) { + std::fill(sw.begin(), sw.end(), 0.0); + std::fill(swI.begin(), swI.end(), 0.0); + for (const int32_t i : sel) { const Obs &o = fulls[i]; - if (!usable(o) || cell[i] < 0 || (parity >= 0 && (o.frame & 1) != parity)) continue; const double a = A[cell[i]], sc = static_cast(o.sigma) * o.corr * a, w = 1.0 / (sc * sc); sw[o.group] += w; swI[o.group] += w * static_cast(o.I) * o.corr * a; } std::vector num(ncell, 0.0), den(ncell, 0.0); - for (size_t i = 0; i < fulls.size(); ++i) { + for (const int32_t i : sel) { const Obs &o = fulls[i]; - if (!usable(o) || cell[i] < 0 || (parity >= 0 && (o.frame & 1) != parity) || sw[o.group] <= 0.0) - continue; + if (sw[o.group] <= 0.0) continue; const double Iref = swI[o.group] / sw[o.group], a = A[cell[i]]; const double Is = static_cast(o.I) * o.corr * a, sc = static_cast(o.sigma) * o.corr * a; if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0) || !(sc > 0.0)) continue; @@ -999,10 +1029,11 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // Mean studentized squared deviation over the subset {frame&1 == parity} with surface A applied, // scored against that subset's own reference (no leakage). Lower = tighter equivalents. auto score = [&](int parity, const std::vector &A) -> double { - std::vector sw(n_groups, 0.0), swI(n_groups, 0.0); - for (size_t i = 0; i < fulls.size(); ++i) { + const std::vector &sel = subset(parity); + std::fill(sw.begin(), sw.end(), 0.0); + std::fill(swI.begin(), swI.end(), 0.0); + for (const int32_t i : sel) { const Obs &o = fulls[i]; - if (!usable(o) || cell[i] < 0 || (o.frame & 1) != parity) continue; const double a = A[cell[i]], Is = static_cast(o.I) * o.corr * a; const double sc = static_cast(o.sigma) * o.corr * a, w = 1.0 / (sc * sc); sw[o.group] += w; swI[o.group] += w * Is; @@ -1012,9 +1043,9 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // reshaping sigma (via corr) without tightening the actual intensities - which on mis-indexed / bad // data passes cross-validation yet worsens Rmeas. A fractional metric cannot be gamed that way. double num = 0.0, den = 0.0; - for (size_t i = 0; i < fulls.size(); ++i) { + for (const int32_t i : sel) { const Obs &o = fulls[i]; - if (!usable(o) || cell[i] < 0 || (o.frame & 1) != parity || sw[o.group] <= 0.0) continue; + if (sw[o.group] <= 0.0) continue; const double a = A[cell[i]], Is = static_cast(o.I) * o.corr * a; const double Iref = swI[o.group] / sw[o.group]; if (!std::isfinite(Iref) || Iref <= 0.0) continue; @@ -1941,7 +1972,9 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool } // French-Wilson amplitudes for IMEAN and (now that they are attached) each Bijvoet hand. - ApplyFrenchWilson(result.merged, x.GetSpaceGroupNumber().value_or(1)); + FrenchWilsonOptions fw_opts; + fw_opts.num_threads = static_cast(nthreads); + ApplyFrenchWilson(result.merged, x.GetSpaceGroupNumber().value_or(1), fw_opts); logger.Info("Merge complete ({} unique reflections)", result.merged.size()); return result; diff --git a/image_analysis/spot_finding/ImageSpotFinder.cpp b/image_analysis/spot_finding/ImageSpotFinder.cpp index 9d0e2939..8296d325 100644 --- a/image_analysis/spot_finding/ImageSpotFinder.cpp +++ b/image_analysis/spot_finding/ImageSpotFinder.cpp @@ -30,24 +30,34 @@ std::vector ImageSpotFinder::Run(const ImagePreprocessorBuffer std::vector ImageSpotFinder::ExtractSpots(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings, const std::vector &res_mask) { - StrongPixelSet pixel_set; - + // Collect the strong pixels first and read their values afterwards, instead of reading the image + // pixel by pixel: on the GPU that read is a device gather, which is what lets the preprocessed + // image stay on the device instead of being copied back in full for every frame. + strong_pixel.clear(); for (int i = 0; i < OutputSize(); i++) { if (output_buffer[i]) { std::bitset<32> bset = output_buffer[i]; for (int bit = 0; bit < 32; bit++) { if (bset.test(bit)) { - size_t npixel = i * 32 + bit; - size_t col = npixel % width; - size_t line = npixel / width; - if (line < height && res_mask[npixel] == 0) - pixel_set.AddStrongPixel(col, line, image[npixel]); + uint32_t npixel = i * 32 + bit; + if (npixel / width < height && res_mask[npixel] == 0) + strong_pixel.push_back(npixel); } } } } std::vector vec; + // The connected-component search rejects a frame with this many strong pixels, so their values are + // of no use. + if (strong_pixel.size() >= UINT16_MAX) + return vec; + + image.Gather(strong_pixel, strong_pixel_value); + + StrongPixelSet pixel_set; + for (size_t i = 0; i < strong_pixel.size(); i++) + pixel_set.AddStrongPixel(strong_pixel[i] % width, strong_pixel[i] / width, strong_pixel_value[i]); pixel_set.FindSpotsImage(settings, vec); return vec; } diff --git a/image_analysis/spot_finding/ImageSpotFinder.h b/image_analysis/spot_finding/ImageSpotFinder.h index b4b938a3..41344c72 100644 --- a/image_analysis/spot_finding/ImageSpotFinder.h +++ b/image_analysis/spot_finding/ImageSpotFinder.h @@ -11,6 +11,10 @@ #include "../image_preprocessing/ImagePreprocessorBuffer.h" class ImageSpotFinder { + // Flat index of every strong pixel of the current image that passed the resolution mask, and its + // value. Kept as members only to reuse the allocation from image to image. + std::vector strong_pixel; + std::vector strong_pixel_value; protected: const int32_t width, height; std::vector output_buffer;