mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 68 Wiki Activity

Find the first pass's spots on every worker

Browse Source 
The first pass of rotation indexing feeds two sampling schemes and a validation set,
and it found their spots one frame at a time on one CUDA stream. It reads 271 to 402
frames depending on the crystal, once per pass, and there are two passes. On a 16M-pixel
dataset that is 6.7 s of a 30 s run at a mean occupancy of 1.9 threads out of 48.

The comment said this was to keep the spot cache from depending on scheduling. The
guarantee is stronger than that and survives: a frame's spots are a pure function of
that frame, so whatever order the workers finish in, the cache ends up holding exactly
what the serial loop put there. The feed order and the point at which the consecutive
scheme stops accumulating are untouched, so the schemes see the same frames as before.

The phase was never CPU-bound, which is why this works: its one thread was spin-waiting
in the driver on the same six full-image passes, about 5.5 ms a frame.

One engine per worker for the whole pass, not per batch. An engine on a 16M-pixel
detector costs ~0.13 s to build - a first attempt rebuilt them per batch and came out
SLOWER than the serial loop, four batches of eight engines against ~1.5 s of frame work.

Two smaller things, in the same per-image path:

The strong-pixel flagging kernel reads four pixels at a time and issues one atomicOr per
four instead of per pixel. Same values, confirmed end to end.

The beam-stop pre-scan sized its per-worker accumulators for the worker count rather than
for the workers that use them. Two int64 arrays per pixel per shard, eight shards: 2.3 GB
allocated and cleared on a 16M-pixel detector for a path the GPU never reads. They are
allocated on use now, which is most of what made that phase look like eight workers with
two of them running.

Measured, three A/B pairs with the order alternated, whole run: 16M-pixel 3600-frame
42.8 s -> 38.8 s, 16M-pixel 1800-frame 38.7 s -> 36.1 s, both at two GPUs; a 4M-pixel
control 23.3 s -> 23.1 s. First-pass indexing itself -44 % and -36 %. The two smaller
datasets merge byte-identically over six runs each; the largest is not bit-reproducible
in the unmodified binary either, and stayed inside that spread.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
jungfrau
2026-08-17 18:06:03 -04:00
co-authored by Claude Opus 5
parent a22a51372e
commit debd73c0bd
4 changed files with 167 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/CHANGELOG.md
+2 -1
View File
@@ -3,7 +3,8 @@
### Unreleased
* **rugnux: scaling and merging are much faster on crystals that integrate far beyond the resolution they merge at.** Observations outside the scaling resolution range are dropped as they are ingested instead of being scaled, combined and post-refined first, and the geometry post-refinement fits on a bounded sample of them. On a large-cell dataset merging at less than half the resolution its detector reaches, a run drops from 68 s to 37 s; merged statistics are unchanged.
* rugnux: the detector-frame modulation correction is fitted on a grid spanning the detector rather than the reflections that happen to be present, so whether it is applied no longer depends on how far integration reached. On a crystal where it was being refused, merged R_meas improves by 4 percentage points.
* rugnux: the beam-stop pre-scan reads and accumulates its frames on several threads instead of one. On a 16M-pixel rotation dataset it drops from 7.9 s to 4.8 s; the shadow it finds is unchanged.
* rugnux: the first-pass rotation indexing finds its spots on every worker rather than one, which is worth most on large detectors - on a 16M-pixel dataset that phase drops by 44% and the whole run by 9%. The lattice it picks is unchanged.
* rugnux: the beam-stop pre-scan reads and accumulates its frames on several threads instead of one. On a 16M-pixel rotation dataset it drops from 7.9 s to 4.8 s; the shadow it finds is unchanged. Its per-worker accumulators are now allocated only where they are used, saving 2.3 GB of host memory on a 16M-pixel detector.
### 1.0.0-rc.161
This is an UNSTABLE release. It includes many experimental features, as well as many AI generated fixes. We recommend using rc.152 for production use.
image_analysis/beam_stop/ShadowFinder.cpp
+12 -8
View File
@@ -302,16 +302,12 @@ ShadowFinder::ShadowFinder(const DiffractionExperiment &experiment, const PixelM
#endif
}
// A shard's accumulators are allocated when a frame is first added to it, not here: with a GPU they
// are never used at all, and on a 16 Mpx detector eight of them are 2.9 GB to allocate and clear -
// which measured 0.8 s of the pre-scan, all of it wasted.
void ShadowFinder::SetShardCount(size_t n) {
const size_t npixels = static_cast<size_t>(width) * height;
shards.clear();
shards.resize(std::max<size_t>(1, n));
for (auto &p : shards) {
p.max_value.assign(npixels, 0);
p.sum_value.assign(npixels, 0);
p.valid_count.assign(npixels, 0);
p.frames = 0;
}
}
ShadowFinder::Projection ShadowFinder::Reduce() const {
@@ -376,7 +372,8 @@ ShadowFinder::Projection ShadowFinder::Reduce() const {
}
}
#endif
for (const auto &p : shards)
for (const auto &p : shards) {
if (p.frames == 0) continue; // never used, and its accumulators were never allocated
for (size_t i = lo; i < hi; i++) {
if (p.valid_count[i] == 0)
continue;
@@ -385,6 +382,7 @@ ShadowFinder::Projection ShadowFinder::Reduce() const {
out.sum_value[i] += p.sum_value[i];
out.valid_count[i] += p.valid_count[i];
}
}
}));
}
for (auto &f : futures) f.get();
@@ -441,6 +439,12 @@ void ShadowFinder::AddImage(const DataMessage &data, std::vector<uint8_t> &buffe
#endif
Projection &p = shards[shard];
if (p.max_value.empty()) {
const size_t npixels = static_cast<size_t>(width) * height;
p.max_value.assign(npixels, 0);
p.sum_value.assign(npixels, 0);
p.valid_count.assign(npixels, 0);
}
const auto ptr = data.image.GetUncompressedPtr(buffer);
switch (data.image.GetMode()) {
case CompressedImageMode::Int8: Add(reinterpret_cast<const int8_t *>(ptr), p); break;
image_analysis/spot_finding/AdaptiveSpotFinderGPU.cu
+32 -1
View File
@@ -230,7 +230,38 @@ __global__ void flag_strong(const int32_t *__restrict__ image,
const float *__restrict__ thr,
uint32_t *__restrict__ strong,
size_t npix, int nbins) {
for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < npix; idx += blockDim.x * gridDim.x) {
// Four pixels per thread, read as one 16-byte and one 8-byte transaction instead of four of
// each, exactly as the ring reduction above reads them - and flagged with a single atomicOr,
// because four consecutive pixels always fall in the same word of the bit buffer. The last
// npix % 4 pixels are done one at a time below, so nothing is read past the end.
const size_t stride = static_cast<size_t>(blockDim.x) * gridDim.x;
const size_t nquad = npix / 4;
for (size_t q = blockIdx.x * blockDim.x + threadIdx.x; q < nquad; q += stride) {
const int4 v4 = reinterpret_cast<const int4 *>(image)[q];
const ushort4 b4 = reinterpret_cast<const ushort4 *>(pixel_to_bin)[q];
const int32_t vq[4] = {v4.x, v4.y, v4.z, v4.w};
const uint16_t bq[4] = {b4.x, b4.y, b4.z, b4.w};
uint32_t bits = 0;
#pragma unroll
for (int k = 0; k < 4; k++) {
const int32_t v = vq[k];
if (v == INT32_MAX) {
bits |= 1u << k;
} else if (v != INT32_MIN) {
const int b = bq[k];
if (b < nbins && static_cast<float>(v) >= thr[b])
bits |= 1u << k;
}
}
if (bits) {
const size_t idx = 4 * q;
atomicOr(&strong[idx / 32], bits << (idx % 32));
}
}
for (size_t idx = 4 * nquad + blockIdx.x * blockDim.x + threadIdx.x; idx < npix; idx += stride) {
const int32_t v = image[idx];
bool s = false;
if (v == INT32_MAX) {
rugnux/Rugnux.cpp
+121 -23
View File
@@ -1177,20 +1177,16 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
};
// Spots for a first-pass image ordinal, cached (a frame is reused across schemes and the
// validation set; re-finding is expensive on the --redo-rotation-spots path). Accessed only
// from single-threaded sections (the scheme feed and the validation loop), so no locking.
// validation set; re-finding is expensive on the --redo-rotation-spots path). The workers below
// hand their spots back and this thread puts them in, so the map itself is still only ever
// touched from here and needs no locking.
std::map<int, std::vector<SpotToSave>> spot_cache;
// One analysis engine for the whole first pass. Constructing it allocates a CUDA stream and a
// full set of GPU buffers (preprocessing, spot finding, azimuthal integration, Bragg
// integration) - far more work than analysing a frame - so building it per image would pay for
// hundreds of them, serially. This section is single-threaded, so one engine is enough.
auto analysis = std::make_unique<MXAnalysisWithoutFPGA>(experiment_, mapping, pixel_mask_, *indexer,
/*enable_fused_adaptive_gpu=*/true);
auto profile = std::make_unique<AzimuthalIntegrationProfile>(mapping);
auto get_spots = [&](int ordinal) -> const std::vector<SpotToSave> & {
auto it = spot_cache.find(ordinal);
if (it != spot_cache.end())
return it->second;
// The spots of one frame, found with the caller's own analysis engine. A frame's spots are a
// pure function of that frame and the settings - the engine carries nothing from one image to
// the next, which is why the main image loop can give one to each of its workers - so it does
// not matter which engine, or which thread, finds them.
const auto find_spots = [&](MXAnalysisWithoutFPGA &analysis, AzimuthalIntegrationProfile &profile,
int ordinal) {
const int image_idx = start_image + ordinal * config_.stride;
std::vector<SpotToSave> spots;
try {
@@ -1204,13 +1200,92 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
m.image = img->image;
if (dataset->efficiency.size() > image_idx)
m.image_collection_efficiency = dataset->efficiency[image_idx];
analysis->Analyze(m, *profile, first_pass);
analysis.Analyze(m, profile, first_pass);
spots = std::move(m.spots);
}
} catch (const std::exception &e) {
logger.Warning("First-pass spot read failed for image {}: {}", image_idx, e.what());
}
return spot_cache.emplace(ordinal, std::move(spots)).first->second;
return spots;
};
// Find the spots of every ordinal in the list that is not cached yet, on several workers, and
// put them in the cache.
//
// The cache ends up holding exactly what a serial loop would have put there: the results are
// collected into a vector indexed by position and inserted afterwards, so neither the spots of
// a frame nor the set of frames in the cache depends on how the workers interleaved.
//
// Worker t keeps its analysis engine for the whole first pass. Building one is a CUDA stream
// and a full set of device buffers - measured at 0.13 s against 4 ms for a frame - and the
// "wedge" scheme only discovers how far into the sweep it has to read as the frames arrive,
// so this is called several times and rebuilding them each time cost a quarter of the pass.
// The engine belongs to a card, not to a thread, so a later call only has to make that card
// current again; workers take a card by index so an engine always meets the one it was
// built on.
const int spot_workers = get_gpu_count() > 0
? std::min(config_.nthreads, std::max(8, 4 * get_gpu_count()))
: config_.nthreads;
struct SpotEngine {
std::unique_ptr<MXAnalysisWithoutFPGA> analysis;
std::unique_ptr<AzimuthalIntegrationProfile> profile;
};
std::vector<SpotEngine> engines(std::max(spot_workers, 1));
const auto use_gpu_of_worker = [](size_t t) {
if (const int32_t gpus = get_gpu_count(); gpus > 0)
set_gpu(static_cast<int32_t>(t) % gpus);
};
int prefetch_batches = 0;
double prefetch_time_s = 0.0;
const auto prefetch_spots = [&](const std::vector<int> &ordinals) {
std::vector<int> wanted;
for (const int ordinal : ordinals)
if (!spot_cache.contains(ordinal))
wanted.push_back(ordinal);
if (wanted.empty())
return;
const auto prefetch_start = std::chrono::steady_clock::now();
++prefetch_batches;
std::vector<std::vector<SpotToSave>> found(wanted.size());
std::atomic<size_t> next{0};
std::vector<std::future<void>> futures;
const size_t nworkers = std::min<size_t>(engines.size(), wanted.size());
futures.reserve(nworkers);
for (size_t t = 0; t < nworkers; t++)
futures.emplace_back(std::async(std::launch::async, [&, t] {
use_gpu_of_worker(t); // must precede engine construction
SpotEngine &e = engines[t];
if (!e.analysis) {
e.analysis = std::make_unique<MXAnalysisWithoutFPGA>(
experiment_, mapping, pixel_mask_, *indexer,
/*enable_fused_adaptive_gpu=*/true);
e.profile = std::make_unique<AzimuthalIntegrationProfile>(mapping);
}
for (size_t i = next.fetch_add(1); i < wanted.size(); i = next.fetch_add(1))
found[i] = find_spots(*e.analysis, *e.profile, wanted[i]);
}));
for (auto &f : futures)
f.get();
for (size_t i = 0; i < wanted.size(); i++)
spot_cache.emplace(wanted[i], std::move(found[i]));
prefetch_time_s += std::chrono::duration<double>(
std::chrono::steady_clock::now() - prefetch_start).count();
};
// Release the engines on workers of their own, so a stream and its device buffers are freed
// with the card they were made on current.
const auto release_spot_engines = [&] {
std::vector<std::future<void>> futures;
futures.reserve(engines.size());
for (size_t t = 0; t < engines.size(); t++)
futures.emplace_back(std::async(std::launch::async, [&, t] {
use_gpu_of_worker(t);
engines[t] = SpotEngine{};
}));
for (auto &f : futures)
f.get();
};
// How many validation frames a candidate global lattice actually indexes per image, scored on
@@ -1227,10 +1302,8 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
idx.ForceRotationIndexerResult(r);
// Fill the spot cache first. It is a plain map filled on demand, and the loop below runs the
// frames in parallel, where a lookup racing an insert is not something the map survives.
// Spot finding stays serial and in frame order, so what ends up in the cache does not depend
// on scheduling; after the first candidate this costs nothing.
for (const int ordinal : validation)
get_spots(ordinal);
// After the first candidate this costs nothing.
prefetch_spots(validation);
// Each frame's verdict is its own - the score is just how many of them index - so they run
// together. Refining one frame's geometry is a few Ceres solves and no GPU, and it is the
// same call the main image loop already makes from every one of its workers on this same
@@ -1251,14 +1324,23 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
};
// Feed one first-pass scheme (a set of image ordinals) into its own rotation indexer, ready to
// be indexed. Spots are pulled from the single-threaded cache here; the FFT + refinement runs
// separately (RunIndexing) so the two schemes' indexing can overlap.
// be indexed. Spots are pulled from the cache here; the FFT + refinement runs separately
// (RunIndexing) so the two schemes' indexing can overlap.
//
// The feed itself stays in frame order and stops where it always did, but it keeps the cache a
// chunk of frames ahead of itself so those frames' spots are found on several workers. A chunk
// rather than the whole list because the "wedge" scheme stops on the accumulation cap, and how
// far into the sweep that is only becomes known as the frames arrive.
constexpr int SPOT_PREFETCH_CHUNK = 64;
auto feed_scheme = [&](IndexerThreadPool &pool, const std::vector<int> &ordinals) {
auto ri = std::make_unique<RotationIndexer>(experiment_, pool);
for (const int ordinal : ordinals) {
for (size_t i = 0; i < ordinals.size(); i++) {
if (cancelled_ || ri->AccumulationFull())
break;
ri->ProcessImage(ordinal, get_spots(ordinal), rot_angle(ordinal));
if (!spot_cache.contains(ordinals[i]))
prefetch_spots({ordinals.begin() + i,
ordinals.begin() + std::min(ordinals.size(), i + SPOT_PREFETCH_CHUNK)});
ri->ProcessImage(ordinals[i], spot_cache.at(ordinals[i]), rot_angle(ordinals[i]));
}
return ri;
};
@@ -1390,6 +1472,19 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
return bp;
};
// Find the spots of everything the first pass is certain to want in one go - both schemes'
// opening frames and the validation set. A batch of workers costs an analysis engine each,
// which is far more than a frame is to analyse, so one batch for the lot beats one per scheme.
{
std::vector<int> upfront = spread_ordinals;
upfront.insert(upfront.end(), validation.begin(), validation.end());
for (int i = 0; i < std::min(SPOT_PREFETCH_CHUNK, images_to_process); i++)
upfront.push_back(i);
std::sort(upfront.begin(), upfront.end());
upfront.erase(std::unique(upfront.begin(), upfront.end()), upfront.end());
prefetch_spots(upfront);
}
FirstPass best = pick_best(*indexer_pool, *indexer);
// Long-axis rescue. When the de-novo cell indexes few validation frames, a long, finely-spaced
@@ -1446,6 +1541,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
best = constrained;
}
}
release_spot_engines(); // no more frames to read; give the cards their memory back
if (!cancelled_) {
// The two ways the first pass can end without a usable lattice. They are one situation from
@@ -1482,6 +1578,8 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
c.a, c.b, c.c, c.alpha, c.beta, c.gamma, next_steps));
}
indexer->ForceRotationIndexerResult(*best.result);
logger.Info("First-pass spot finding: {} frames in {} batches, {:.2f} s",
spot_cache.size(), prefetch_batches, prefetch_time_s);
logger.Info("Two-pass rotation indexing found lattice (scheme '{}': {}/{} validation frames)",
best.name, best.score, static_cast<int>(validation.size()));