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. A worker reserves its bytes before it makes them and hands
// them back when it is done; one that would take the pool past its size waits until others have handed
// theirs back - except that a reservation made while nothing is reserved 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);
    [[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;
    };

private:
    HostWorkerPool();
    std::mutex m;
    std::condition_variable cv;
    size_t size = 0;
    size_t reserved = 0;
};