partition.hpp

Go to the documentation of this file.

#pragma once

 
#include <cstddef>
 
#include "orrery/core/aligned_allocator.hpp"
#include "orrery/core/types.hpp"
 
namespace orrery::backend {

inline constexpr core::Index kPartitionGrain = core::kCacheLineBytes / sizeof(core::Real);

struct IndexRange {
    core::Index begin{};
    core::Index end{};
 
    [[nodiscard]] constexpr bool empty() const noexcept { return begin >= end; }
 
    [[nodiscard]] constexpr core::Index size() const noexcept { return empty() ? 0 : end - begin; }
};

[[nodiscard]] constexpr IndexRange equal_share(core::Index count, unsigned workers,
                                               unsigned worker) noexcept {
    if (workers == 0 || worker >= workers) {
        return IndexRange{count, count};
    }
 
    const core::Index lines = (count + kPartitionGrain - 1) / kPartitionGrain;
    const core::Index share = lines / workers;
    const core::Index remainder = lines % workers;
    const core::Index index = worker;
 
    const core::Index first_line = (index * share) + (index < remainder ? index : remainder);
    const core::Index line_count = share + (index < remainder ? 1 : 0);
 
    const core::Index begin = first_line * kPartitionGrain;
    const core::Index end = (first_line + line_count) * kPartitionGrain;
 
    // The last worker's range runs off the end of the array, because the count
    // is not generally a whole number of cache lines. Clamping here rather than
    // in every caller keeps the boundary arithmetic above in one piece.
    return IndexRange{begin < count ? begin : count, end < count ? end : count};
}
 
} // namespace orrery::backend