Jungfraujoch/common/HostMemory.h

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

#pragma once

#include <condition_variable>
#include <cstddef>
#include <mutex>

// The machine's physical memory in bytes - its total, not what is free at the moment. 0 where it
// cannot be read.
size_t get_host_total_memory();

// The host memory a run plans against: this fraction of the machine's physical memory, decided from
// the total, never from what happens to be free (as what runs on a GPU is decided on the card's total).
// The plan decides only what runs at the same time - two merges side by side or one after the other,
// how many large per-worker buffers are held at once - and never what is computed, so no result
// depends on it.
constexpr double HOST_MEMORY_BUDGET_FRACTION = 0.75;
size_t DefaultHostMemoryBudget();

// The share of the budget that the large per-worker buffers of parallel steps - a translation function's
// grids, a null replicate's model maps - may hold at once. The rest is what the run holds beside them:
// its reflections, its merge, the detector-sized arrays.
constexpr double HOST_WORKER_POOL_FRACTION = 0.5;

// One pool for those buffers, process-wide, held in one ledger with what the run holds beside them. The run
// says what it holds - its reflections, its merge engine, a validation's own grids - as Held bytes, and the
// workers may hold at once the pool's size or what the budget leaves beside the Held bytes, whichever is
// less: so the two together never claim more than the budget, however large the run's own part is. A worker
// reserves its bytes before it makes them and hands them back when it is done; one that would take the pool
// past that waits until others have handed theirs back - except that a reservation made while no worker holds
// one is always let through, however large. So how many such workers run at once is decided by their bytes,
// not by the thread count, and however the steps holding them nest or run side by side; what each worker
// computes, and where its result goes, does not change. A worker must not reserve again while it holds a
// reservation, and none may reserve on ParallelFor's pool: a holder's own parallel passes need those threads.
// Work spread over that pool is bounded by its count instead (Size() / bytes each).
class HostWorkerPool {
public:
    static HostWorkerPool &Instance();

    // The pool's size in bytes; 0 sets it back to HOST_WORKER_POOL_FRACTION of DefaultHostMemoryBudget().
    void SetSize(size_t bytes);
    // The budget the pool and the Held bytes share; 0 sets it back to DefaultHostMemoryBudget().
    void SetBudget(size_t bytes);
    // What the workers may hold at once now: the pool's size, or the budget less the Held bytes if less.
    [[nodiscard]] size_t Size();

    class Reservation {
    public:
        Reservation(HostWorkerPool &pool, size_t bytes);
        ~Reservation();
        Reservation(const Reservation &) = delete;
        Reservation &operator=(const Reservation &) = delete;
    private:
        HostWorkerPool &pool;
        size_t bytes;
    };

    // Bytes the run holds beside the workers for as long as this lives; never waits.
    class Held {
    public:
        Held(HostWorkerPool &pool, size_t bytes);
        ~Held();
        Held(const Held &) = delete;
        Held &operator=(const Held &) = delete;
        void Set(size_t bytes);   // what it holds from now on
    private:
        HostWorkerPool &pool;
        size_t bytes;
    };

private:
    HostWorkerPool();
    size_t available() const;   // Size(), with the lock held
    std::mutex m;
    std::condition_variable cv;
    size_t size = 0;
    size_t budget = 0;
    size_t held = 0;
    size_t reserved = 0;
};