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 "../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

void WarmUpCuFFT() {
#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() == 0)
        return;
    cufftHandle plan = 0;
    if (cufftPlan1d(&plan, 1024, CUFFT_C2C, 1) == CUFFT_SUCCESS)
        cufftDestroy(plan);
#endif
}

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

FFTPlanReuse::~FFTPlanReuse() = default;

#ifdef JFJOCH_USE_CUDA
// 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) {
    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);
}
#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))
                return;
            if (state == TaskState::READY) {
                state = TaskState::RUNNING;
                input = std::move(task_input);
            } else
                drop = true;
        }
        if (drop) {
            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.
    for (int32_t d = 0; d < get_gpu_count(); d++) {
        ScopedGpuPin pin(d);
        cuda_trim_allocation_pool();
    }
#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};
}