Jungfraujoch/image_analysis/indexing/IndexerThreadPool.cpp

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

#include "IndexerThreadPool.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <map>
#include <tuple>

#include "../common/CUDAWrapper.h"
#include "../common/Logger.h"

#ifdef JFJOCH_USE_CUDA
#include "FFBIDXIndexer.h"
#include "FFTIndexerGPU.h"
#endif

#ifdef JFJOCH_USE_FFTW
#include "FFTIndexerCPU.h"
#endif

double WarmUpCuFFT() {
#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() == 0)
        return NAN;
    // The first plan pays the initialisation, and compiles its kernels if the JIT cache does not have them
    // yet; the next ones show what a plan costs from then on. The fastest of three, so that a stall
    // beside (another card's context being made) is not taken for a slow plan.
    double fastest_ms = INFINITY;
    for (int i = 0; i < 4; i++) {
        const auto t0 = std::chrono::steady_clock::now();
        cufftHandle plan = 0;
        if (cufftPlan3d(&plan, 64, 64, 64, CUFFT_R2C) != CUFFT_SUCCESS)
            return NAN;
        cufftDestroy(plan);
        const double ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count();
        if (i > 0)
            fastest_ms = std::min(fastest_ms, ms);
    }
    return fastest_ms;
#else
    return NAN;
#endif
}

#ifdef JFJOCH_USE_CUDA
FFTPlanReuse::FFTPlanReuse() : scope(std::make_unique<CudaFFTPlanScope>()) {}
#else
class CudaFFTPlanScope {};
FFTPlanReuse::FFTPlanReuse() = default;
#endif

FFTPlanReuse::~FFTPlanReuse() = default;

#ifndef JFJOCH_USE_CUDA
FFTIndexerReuse::FFTIndexerReuse() = default;
FFTIndexerReuse::~FFTIndexerReuse() = default;
void KeepFFTIndexers(const IndexingSettings &, int) {}
#endif

#ifdef JFJOCH_USE_CUDA
namespace {
    // The card and what FFTIndexer's constructor reads from the settings.
    using FFTIndexerKey = std::tuple<int, float, float, float, float, int64_t, float, int64_t>;

    struct KeptFFTIndexers {
        std::mutex m;
        int scopes = 0;
        std::map<FFTIndexerKey, std::vector<std::unique_ptr<Indexer>>> idle;
    };

    KeptFFTIndexers &kept_fft_indexers() {
        static auto *k = new KeptFFTIndexers;
        return *k;
    }

    FFTIndexerKey KeyOnCurrentCard(const IndexingSettings &s) {
        int device = 0;
        if (cudaGetDevice(&device) != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Failed to read the CUDA device");
        return {device, s.GetFFT_MaxUnitCell_A(), s.GetFFT_MinUnitCell_A(), s.GetFFT_MinAngle_deg(),
                s.GetFFT_MaxAngle_deg(), s.GetFFT_NumVectors(), s.GetFFT_HighResolution_A(), s.GetRefineThreads()};
    }

    std::unique_ptr<Indexer> BorrowFFTIndexer(const IndexingSettings &settings) {
        auto &k = kept_fft_indexers();
        std::lock_guard lock(k.m);
        auto it = k.idle.find(KeyOnCurrentCard(settings));
        if (it == k.idle.end() || it->second.empty())
            return nullptr;
        auto ret = std::move(it->second.back());
        it->second.pop_back();
        return ret;
    }

    // On the card the indexer was built on, which is the calling worker's.
    void GiveBackFFTIndexer(const IndexingSettings &settings, std::unique_ptr<Indexer> &indexer) {
        if (!indexer)
            return;
        auto &k = kept_fft_indexers();
        std::lock_guard lock(k.m);
        if (k.scopes > 0)
            k.idle[KeyOnCurrentCard(settings)].push_back(std::move(indexer));
    }
}

FFTIndexerReuse::FFTIndexerReuse() {
    std::lock_guard lock(kept_fft_indexers().m);
    kept_fft_indexers().scopes++;
}

FFTIndexerReuse::~FFTIndexerReuse() {
    std::map<FFTIndexerKey, std::vector<std::unique_ptr<Indexer>>> idle;
    {
        std::lock_guard lock(kept_fft_indexers().m);
        if (--kept_fft_indexers().scopes == 0)
            idle.swap(kept_fft_indexers().idle);
    }
    for (auto &[key, list] : idle) {
        ScopedGpuPin pin(std::get<0>(key));
        list.clear();
    }
}

// An FFT indexer is a cuFFT plan and ~0.6-0.9 GB of device buffers, and several pipeline copies of one
// process run on the card at the same time: probe passes (GPUWorkBeside) beside the main line, for a
// few seconds each. So one that does not fit hands back what the allocation pool holds unused and waits
// for the work beside it to end, then is built again - as the merge's allocations do. The main line
// waits for the probes; a probe waits for the main line's image pass (GPUImagePass), never for a probe,
// itself included. Neither of those waits for the other, so nothing waits in a circle. Only when the
// indexer is built depends on this, never what it computes; one that does not fit the card fails as
// before.
static std::unique_ptr<Indexer> MakeFFTIndexerGPU(const IndexingSettings &settings, bool beside) {
    if (auto kept = BorrowFFTIndexer(settings))
        return kept;
    try {
        return std::make_unique<FFTIndexerGPU>(settings);
    } catch (const JFJochException &) {
        cuda_throw_if_context_lost();   // not memory: waiting would only hide the real error
        cuda_clear_error();
    }
    if (beside)
        wait_for_gpu_image_passes();
    else
        // That work is a probe pass of seconds, so this long means it is stuck: the indexer fails as
        // it would have without it.
        wait_for_gpu_work_beside(std::chrono::minutes(10));
    cuda_trim_allocation_pool();
    return std::make_unique<FFTIndexerGPU>(settings);
}

void KeepFFTIndexers(const IndexingSettings &settings, int count) {
    if (get_gpu_count() == 0)
        return;
    std::vector<std::unique_ptr<Indexer>> indexers;
    for (int i = 0; i < count; i++)
        indexers.push_back(MakeFFTIndexerGPU(settings, false));
    for (auto &i : indexers)
        GiveBackFFTIndexer(settings, i);
}
#endif

// The indexer for one RESOLVED algorithm, or nullptr if this build/host cannot serve it.
static std::unique_ptr<Indexer> MakeIndexer(IndexingAlgorithmEnum algorithm, const IndexingSettings &settings,
                                            bool beside) {
#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() > 0) {
        if (algorithm == IndexingAlgorithmEnum::FFT)
            return MakeFFTIndexerGPU(settings, beside);
        if (algorithm == IndexingAlgorithmEnum::FFBIDX)
            return std::make_unique<FFBIDXIndexer>();
    }
#endif
#ifdef JFJOCH_USE_FFTW
    if (algorithm == IndexingAlgorithmEnum::FFTW)
        return std::make_unique<FFTIndexerCPU>(settings);
#endif
    return nullptr;
}

IndexerThread::IndexerThread(const IndexingSettings &settings, int threadid, IndexerConstruction construction,
                             bool beside)
    : settings_(settings), construction_(construction), beside_(beside) {
    std::unique_lock<std::mutex> lock(m);
    state = TaskState::STARTING;
    worker_thread = std::thread(&IndexerThread::Worker, this, threadid);
    c_running.wait(lock, [this] { return state != TaskState::STARTING; });
    if (state == TaskState::ERROR) {
        worker_thread.join();
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Indexer thread initialization failed");
    }
}

void IndexerThread::Worker(int threadid) {
    try {
        pin_gpu();
    } catch (const std::exception &e) {
        spdlog::error("Failed to pin to GPU {}", e.what());
    } catch (...) {
        // GPU pinning errors are not critical and should be ignored for the time being.
    }

    std::unique_ptr<Indexer> fft_indexer, ffbidx_indexer, fftw_indexer;

    // Preconstruct: build every indexer the requested algorithm could resolve to before the pool
    // reports ready, so no cuFFT planning happens once frames are flowing, and a failure is fatal
    // for the pool instead of being met frame by frame. OnFirstUse skips this and builds in the
    // dispatch below.
    if (construction_ == IndexerConstruction::Preconstruct) {
        try {
            const auto requested = settings_.GetAlgorithm();

            if (requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFT)
                fft_indexer = MakeIndexer(IndexingAlgorithmEnum::FFT, settings_, beside_);

            if (requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFBIDX)
                ffbidx_indexer = MakeIndexer(IndexingAlgorithmEnum::FFBIDX, settings_, beside_);

            if ((requested == IndexingAlgorithmEnum::Auto && get_gpu_count() == 0)
                || requested == IndexingAlgorithmEnum::FFTW)
                fftw_indexer = MakeIndexer(IndexingAlgorithmEnum::FFTW, settings_, beside_);
        } catch (const std::exception &e) {
            spdlog::error("Failed to initialize indexer: {}", e.what());
            {
                std::unique_lock<std::mutex> lock(m);
                state = TaskState::ERROR;
            }
            c_running.notify_all();
            return;
        } catch (...) {
            spdlog::error("Failed to initialize indexer");
            {
                std::unique_lock<std::mutex> lock(m);
                state = TaskState::ERROR;
            }
            c_running.notify_all();
            return;
        }
    }

    {
        std::unique_lock<std::mutex> lock(m);
        state = TaskState::IDLE;
    }
    c_running.notify_all();

    while (true) {
        std::unique_ptr<TaskInput> input;
        bool drop = false;
        // Look for task + handle stop
        {
            std::unique_lock<std::mutex> lock(m);
            c_start.wait(lock, [this] { return stop || drop_indexers || state == TaskState::READY; });
            if (stop && (state != TaskState::READY)) {
#ifdef JFJOCH_USE_CUDA
                GiveBackFFTIndexer(settings_, fft_indexer);
#endif
                return;
            }
            if (state == TaskState::READY) {
                state = TaskState::RUNNING;
                input = std::move(task_input);
            } else
                drop = true;
        }
        if (drop) {
#ifdef JFJOCH_USE_CUDA
            GiveBackFFTIndexer(settings_, fft_indexer);
#endif
            fft_indexer.reset();
            ffbidx_indexer.reset();
            fftw_indexer.reset();
            {
                std::unique_lock<std::mutex> lock(m);
                drop_indexers = false;
            }
            c_done.notify_all();
        }
        if (input) {
            std::unique_ptr<IndexerResult> tmp_result;
            try {
                auto algorithm = input->experiment.GetIndexingAlgorithm();
                std::unique_ptr<Indexer> *slot = nullptr;

                switch (algorithm) {
                    case IndexingAlgorithmEnum::FFT:    slot = &fft_indexer;    break;
                    case IndexingAlgorithmEnum::FFBIDX: slot = &ffbidx_indexer; break;
                    case IndexingAlgorithmEnum::FFTW:   slot = &fftw_indexer;   break;
                    default: break;
                }

                // A preconstructing worker already holds it; an OnFirstUse worker builds it here,
                // on the first frame that resolves to this algorithm.
                if (slot && !*slot)
                    *slot = MakeIndexer(algorithm, settings_, beside_);

                if (!slot || !*slot) {
                    // Algorithm is already resolved here (never Auto/None - see
                    // IndexerThreadPool::Run, which also checked this host can serve it). Reaching
                    // this means the resolved algorithm has no matching indexer in this build -
                    // fail loudly instead of silently not indexing.
                    throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                          "Internal error: no indexer available for the resolved "
                                          "indexing algorithm");
                }

                Indexer &indexer = **slot;
                indexer.Setup(input->experiment);
                tmp_result = std::make_unique<IndexerResult>(indexer.Run(input->recip, input->severity_only));
            } catch (std::exception &e) {
                // Hand the failure back as a result carrying the reason. A nullptr here was
                // indistinguishable from a worker that was never dispatched, and both then read
                // downstream as "this frame did not index".
                spdlog::error("Indexer thread {} failed: {}", threadid, e.what());
                tmp_result = std::make_unique<IndexerResult>(IndexerResult{
                        .lattice = {}, .indexing_time_s = 0, .executed = false,
                        .error = std::current_exception()});
                // This thread goes on to the next frame, so a CUDA error left behind here would be
                // reported by the first cudaGetLastError() of every frame that follows - one failed
                // attempt would read as an indexer that never works again.
                cuda_clear_error();
            }
            {
                std::unique_lock<std::mutex> lock(m);
                state = TaskState::COMPLETED;
                result = std::move(tmp_result);
            }
            c_done.notify_all();
        }
    }
}

void IndexerThread::Release() {
    std::unique_lock<std::mutex> lock(m);
    if (stop)
        return;
    drop_indexers = true;
    c_start.notify_all();
    c_done.wait(lock, [this] { return !drop_indexers; });
}

void IndexerThread::Finalize() {
    {
        std::unique_lock<std::mutex> lock(m);
        stop = true;
    }
    c_start.notify_all();
    if (worker_thread.joinable())
        worker_thread.join();
}

std::unique_ptr<IndexerResult> IndexerThread::Run(const DiffractionExperiment &experiment,
                                                  const std::vector<Coord> &recip, bool severity_only) {
    std::unique_ptr<IndexerResult> tmp_result;
    {
        std::unique_lock<std::mutex> lock(m);
        if (stop)
            return nullptr;
        if (state != TaskState::IDLE)
            return nullptr;
        task_input = std::make_unique<TaskInput>(std::cref(experiment), std::cref(recip), severity_only);
        state = TaskState::READY;
    }
    c_start.notify_one();
    {
        std::unique_lock<std::mutex> lock(m);
        c_done.wait(lock, [this] { return state == TaskState::COMPLETED; });
        tmp_result = std::move(result);
        state = TaskState::IDLE;
    }
    return tmp_result;
}

IndexerThread::~IndexerThread() {
    Finalize();
}

IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings, IndexerConstruction construction,
                                     bool beside)
    : worker_busy(settings.GetIndexingThreads(), 0),
      worker_free_count(settings.GetIndexingThreads()),
      viable_cell_min_spots(settings.GetViableCellMinSpots()),
      blocking(settings.GetBlockingBehavior()),
      settings_(settings) {
    for (size_t i = 0; i < settings.GetIndexingThreads(); ++i)
        tasks.emplace_back(std::make_unique<IndexerThread>(std::cref(settings), i, construction, beside));
}

void IndexerThreadPool::ReleaseIndexers() {
    for (auto &t : tasks)
        t->Release();
#ifdef JFJOCH_USE_CUDA
    // Their buffers went back to the allocation pool, which would keep them reserved for the next
    // buffer of whoever asks; hand them to the card.
    cuda_trim_allocation_pools();
#endif
}

int IndexerThreadPool::GetFreeWorker() {
    std::unique_lock<std::mutex> lock(m);

    if (tasks.size() == 0)
        return -1;

    if (blocking)
        c.wait(lock, [this] { return worker_free_count > 0; });

    for (int i = 0; i < tasks.size(); i++) {
        if (worker_busy[i] == 0) {
            worker_busy[i] = 1;
            worker_free_count--;
            return i;
        }
    }
    return -1;
}

IndexerResult IndexerThreadPool::Run(const DiffractionExperiment &experiment, const std::vector<Coord> &recip,
                                     bool severity_only) {
    const auto algorithm = experiment.GetIndexingAlgorithm();
    if (algorithm == IndexingAlgorithmEnum::None)
        return IndexerResult{.lattice = {}, .indexing_time_s = 0, .executed = false};

    // GetIndexingAlgorithm() must already have resolved Auto to a concrete algorithm;
    // the pool has no policy to resolve it, so Auto here is an upstream contract bug.
    if (algorithm == IndexingAlgorithmEnum::Auto)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Internal error: indexing algorithm must be resolved (not Auto) "
                              "before reaching the indexer pool");

    // The workers built their indexers from the raw requested algorithm, but the algorithm actually
    // dispatched is the RESOLVED one (rotation, for instance, always resolves to the GPU FFT indexer
    // when a GPU is present, ignoring the request). If the resolution lands on an algorithm this host
    // did not build an indexer for, fail here with an explanation instead of the opaque "no indexer
    // available for the resolved algorithm" from deep inside a worker.
    const auto requested = experiment.GetIndexingSettings().GetAlgorithm();
    const bool have_gpu = get_gpu_count() > 0;
#ifdef JFJOCH_USE_FFTW
    constexpr bool fftw_built = true;
#else
    constexpr bool fftw_built = false;
#endif
    const bool servable =
        (algorithm == IndexingAlgorithmEnum::FFT    && have_gpu &&
             (requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFT)) ||
        (algorithm == IndexingAlgorithmEnum::FFBIDX && have_gpu &&
             (requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFBIDX)) ||
        (algorithm == IndexingAlgorithmEnum::FFTW   && fftw_built &&
             ((requested == IndexingAlgorithmEnum::Auto && !have_gpu) || requested == IndexingAlgorithmEnum::FFTW));
    if (!servable) {
        std::string msg;
        if (requested == IndexingAlgorithmEnum::FFTW && have_gpu)
            msg = "FFTW is the CPU indexer and is not available on a node with a GPU. Rotation indexing "
                  "always uses the GPU FFT indexer here; select FFT or Auto, or run FFTW on a CPU-only node.";
        else if (algorithm == IndexingAlgorithmEnum::FFT && !have_gpu)
            msg = "FFT is the GPU indexer but no GPU is available. Select FFTW or Auto for CPU indexing.";
        else if (algorithm == IndexingAlgorithmEnum::FFBIDX && !have_gpu)
            msg = "FFBIDX is a GPU indexer but no GPU is available. Select FFTW or Auto for CPU indexing.";
        else if (algorithm == IndexingAlgorithmEnum::FFTW)
            msg = "FFTW (CPU) indexing was requested but this build has no FFTW indexer.";
        else
            msg = "the requested indexing algorithm resolved to one with no indexer available on this host.";
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot index: " + msg);
    }

    // Check if there is available worker
   const int task = GetFreeWorker();

    std::unique_ptr<IndexerResult> result;
    if (task >= 0) {
        try {
            result = tasks[task]->Run(experiment, recip, severity_only);
        } catch (const std::exception &e) {
            spdlog::error("Indexer thread failed: {}", e.what());
            result = std::make_unique<IndexerResult>(IndexerResult{
                    .lattice = {}, .indexing_time_s = 0, .executed = false,
                    .error = std::current_exception()});
        }
        {
            std::unique_lock<std::mutex> lock(m);
            worker_busy[task] = 0;
            worker_free_count++;
        }
        c.notify_one();
    }
    if (result)
        return *result;
    // No free worker, or the pool is stopping: indexing was not attempted. Distinct from both a
    // frame that did not index and an indexer that failed, and left without an error for that reason.
    return IndexerResult{.lattice = {}, .indexing_time_s = 0};
}