Jungfraujoch/image_analysis/indexing/IndexerThreadPool.h

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once


#include <thread>
#include <mutex>
#include <condition_variable>
#include <queue>
#include <functional>
#include <future>
#include <vector>
#include <optional>
#include <memory>
#include <latch>

#include "../common/JFJochMessages.h"
#include "../common/DiffractionSpot.h"
#include "../common/DiffractionExperiment.h"
#include "Indexer.h"

// When a worker builds its indexers.
//
// Preconstruct (the default, and what the online service needs): every indexer the requested
// algorithm could resolve to is built while the pool is being constructed, so it is resident and
// planned before the first frame arrives. Building one means a cuFFT plan plus ~0.6 GB of device
// allocation, and the first cuFFT call in a process also pays the library's one-time init (0.3 s
// measured); jfjoch_broker cannot have any of that land on the live data path, so holding an indexer
// the resolved algorithm may never dispatch to is an accepted cost there.
//
// OnFirstUse (offline batch - rugnux, jfjoch_viewer): nothing is latency-critical, so the indexer is
// built on the first frame that needs it. The algorithm is only resolved from the frame's
// DiffractionExperiment (rotation -> FFT, stills with a known cell -> FFBIDX), so preconstructing
// leaves each worker holding a fully allocated FFTIndexerGPU - 0.6 GB of cuFFT plan and histograms -
// that can never be dispatched to.
enum class IndexerConstruction { Preconstruct, OnFirstUse };

// Force cuFFT's one-time library initialisation - 0.3 s measured, paid by whichever call happens to
// be first - so it can be got out of the way on a background thread at startup rather than landing in
// the middle of the first pass with nothing overlapping it. Returns how long a 3D plan takes to make
// after that, in ms - about 1 from a working CUDA JIT cache, 60 or more without (see CudaJitCache.h) -
// and NaN without a GPU.
double WarmUpCuFFT();

// cuFFT plans are kept and reused for as long as one of these lives, and destroyed when the last one
// ends (CudaFFTPlanScope): one per processing run in `rugnux`, so no plan outlives the run that made it;
// one for its lifetime in jfjoch_broker. Nothing without CUDA.
class CudaFFTPlanScope;
class FFTPlanReuse {
    std::unique_ptr<CudaFFTPlanScope> scope;
public:
    FFTPlanReuse();
    ~FFTPlanReuse();
    FFTPlanReuse(const FFTPlanReuse &) = delete;
    FFTPlanReuse &operator=(const FFTPlanReuse &) = delete;
};

// GPU FFT indexers are kept for reuse while one of these is alive, by card and the settings they are
// built from, and destroyed when the last one ends. A pool's worker then borrows a kept indexer instead
// of building one, and gives it back - not to the card - when it ends or is released (ReleaseIndexers):
// the device memory stays taken in between. For jfjoch_broker's rugnux live runs, which build the same
// pools for every pass of every collection; an indexer computes the same with any earlier frames behind
// it, as it does within a pool. Nothing without CUDA.
class FFTIndexerReuse {
public:
    FFTIndexerReuse();
    ~FFTIndexerReuse();
    FFTIndexerReuse(const FFTIndexerReuse &) = delete;
    FFTIndexerReuse &operator=(const FFTIndexerReuse &) = delete;
};

// Builds `count` GPU FFT indexers of these settings on the current card at the same time, and keeps them
// for the pools that ask for them later (inside an FFTIndexerReuse; outside one they are only built and
// destroyed). Nothing without a GPU.
void KeepFFTIndexers(const IndexingSettings &settings, int count);

class IndexerThread {
    struct TaskInput {
        const DiffractionExperiment &experiment;
        const std::vector<Coord> &recip;
        const bool severity_only;
    };

    // Held by value: with IndexerConstruction::OnFirstUse the worker builds its indexer long after
    // the constructor returned, and pools are routinely built from a temporary - for instance
    // IndexerThreadPool(experiment.GetIndexingSettings()), which returns by value.
    const IndexingSettings settings_;
    const IndexerConstruction construction_;
    const bool beside_;

    bool stop = false;
    bool drop_indexers = false;   // see Release()
    enum class TaskState {STARTING, IDLE, READY, RUNNING, COMPLETED, ERROR} state = TaskState::STARTING;
    std::mutex m;
    std::condition_variable c_running;
    std::condition_variable c_start;
    std::condition_variable c_done;
    std::unique_ptr<IndexerResult> result = nullptr;
    std::unique_ptr<TaskInput> task_input = nullptr;
    std::thread worker_thread;

    void Worker(int threadid);
public:
    IndexerThread(const IndexingSettings& settings, int threadid, IndexerConstruction construction, bool beside);
    ~IndexerThread();
    std::unique_ptr<IndexerResult> Run(const DiffractionExperiment &experiment, const std::vector<Coord> &recip,
                                       bool severity_only = false);
    // The thread drops the indexers it has built - once the frame it is on, if any, is done - and builds
    // them again on the next frame that needs one. Returns when they are gone.
    void Release();
    void Finalize();
};

class IndexerThreadPool {
    std::mutex m;
    std::condition_variable c;
    std::vector<uint8_t> worker_busy;
    size_t worker_free_count;
    std::vector<std::unique_ptr<IndexerThread>> tasks;
    const int64_t viable_cell_min_spots;
    const bool blocking;
    const IndexingSettings settings_;
    int GetFreeWorker();
public:
    // beside: the pool serves a probe pass that holds GPUWorkBeside, so an FFT indexer that does not fit
    // the card waits for the main line's image pass to end, not for the probes - it is one of them (see
    // MakeFFTIndexerGPU).
    IndexerThreadPool(const IndexingSettings& settings,
                      IndexerConstruction construction = IndexerConstruction::Preconstruct,
                      bool beside = false);
    // The settings the pool's indexers were built with.
    [[nodiscard]] const IndexingSettings &Settings() const { return settings_; }
    // severity_only skips indexing and produces just the spindle severity - see Indexer::Run.
    IndexerResult Run(const DiffractionExperiment& experiment, const std::vector<Coord>& recip,
                      bool severity_only = false);
    // Give back to the card the device memory of the indexers built so far - an FFT indexer holds
    // 0.6-1.2 GB of it - for a stretch with nothing to index. A frame that comes after all builds its
    // indexer again, as on first use, and the indexer it gets computes what the dropped one would.
    void ReleaseIndexers();
};