numeric.hpp

Go to the documentation of this file.

// -------------------------------------------------------//
//
// SHAMROCK code for hydrodynamics
// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
//
// -------------------------------------------------------//
 
#pragma once
 
#include "shambase/assert.hpp"
#include "shambase/memory.hpp"
#include "shamalgs/collective/reduction.hpp"
#include "shambackends/DeviceBuffer.hpp"
#include "shambackends/DeviceScheduler.hpp"
#include "shambackends/kernel_call.hpp"
#include "shambackends/sycl.hpp"
 
namespace shamalgs::numeric {
 
    template<class T>
    sycl::buffer<T> scan_exclusive(sycl::queue &q, sycl::buffer<T> &buf1, u32 len);
 
    template<class T>
    sham::DeviceBuffer<T> scan_exclusive(
        sham::DeviceScheduler_ptr sched, sham::DeviceBuffer<T> &buf1, u32 len);
 
    template<class T>
    sycl::buffer<T> scan_inclusive(sycl::queue &q, sycl::buffer<T> &buf1, u32 len);
 
    template<class T>
    void scan_exclusive_in_place(sycl::queue &q, sycl::buffer<T> &buf, u32 len);
 
    template<class T>
    void scan_inclusive_in_place(sycl::queue &q, sycl::buffer<T> &buf, u32 len);
 
    std::tuple<std::optional<sycl::buffer<u32>>, u32> stream_compact(
        sycl::queue &q, sycl::buffer<u32> &buf_flags, u32 len);
 
    sham::DeviceBuffer<u32> stream_compact(
        const sham::DeviceScheduler_ptr &sched, sham::DeviceBuffer<u32> &buf_flags, u32 len);
 
    template<class T>
    struct histogram_result {
        sham::DeviceBuffer<u64> counts;
        sham::DeviceBuffer<T> bins_center;
        sham::DeviceBuffer<T> bins_width;
    };
 
    template<class Tret, class T>
    sham::DeviceBuffer<Tret> device_histogram(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges,
        u64 nbins,
        const sham::DeviceBuffer<T> &values,
        u32 len);
 
    template<class T>
    inline sham::DeviceBuffer<u64> device_histogram_u64(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges,
        u64 nbins,
        const sham::DeviceBuffer<T> &values,
        u32 len) {
        return device_histogram<u64, T>(sched, bin_edges, nbins, values, len);
    }
 
    template<class T>
    inline sham::DeviceBuffer<u32> device_histogram_u32(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges,
        u64 nbins,
        const sham::DeviceBuffer<T> &values,
        u32 len) {
        return device_histogram<u32, T>(sched, bin_edges, nbins, values, len);
    }
 
    template<class T>
    histogram_result<T> device_histogram_full(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges,
        u64 nbins,
        const sham::DeviceBuffer<T> &values,
        u32 len) {
 
        SHAM_ASSERT(nbins > 1); // at least a sup and a inf
        SHAM_ASSERT(bin_edges.get_size() == nbins + 1);
 
        auto &q = shambase::get_check_ref(sched).get_queue();
 
        sham::DeviceBuffer<u64> counts = device_histogram(sched, bin_edges, nbins, values, len);
 
        sham::DeviceBuffer<T> bins_center(nbins, sched);
        sham::DeviceBuffer<T> bins_width(nbins, sched);
 
        sham::kernel_call(
            q,
            sham::MultiRef{bin_edges},
            sham::MultiRef{bins_center, bins_width},
            nbins,
            [](u32 i,
               const T *__restrict bin_edges,
               T *__restrict bins_center,
               T *__restrict bins_width) {
                bins_center[i] = (bin_edges[i] + bin_edges[i + 1]) / 2;
                bins_width[i]  = bin_edges[i + 1] - bin_edges[i];
            });
 
        return {std::move(counts), std::move(bins_center), std::move(bins_width)};
    }
 
    template<class T>
    struct BinnedCompute {
        sham::DeviceBuffer<T>
            valid_values; 
        sham::DeviceBuffer<u32> offsets_bins; 
    };
 
    template<class T>
    BinnedCompute<T> binned_init_compute(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges,
        u64 nbins,
        const sham::DeviceBuffer<T> &values, // ie f(r)
        const sham::DeviceBuffer<T> &keys,   // ie r
        u32 len);
 
    template<class Tret, class T, class Fct>
    sham::DeviceBuffer<Tret> binned_compute(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges,
        u64 nbins,
        const sham::DeviceBuffer<T> &values, // ie f(r)
        const sham::DeviceBuffer<T> &keys,   // ie r
        u32 len,
        Fct &&fct) {
 
        auto bin_compute_in = binned_init_compute(sched, bin_edges, nbins, values, keys, len);
        auto &valid_values  = bin_compute_in.valid_values;
        auto &offsets_bins  = bin_compute_in.offsets_bins;
 
        auto &q = shambase::get_check_ref(sched).get_queue();
 
        sham::DeviceBuffer<Tret> bin_compute(nbins, sched);
 
        sham::kernel_call(
            q,
            sham::MultiRef{valid_values, offsets_bins},
            sham::MultiRef{bin_compute},
            nbins,
            [fct](
                u32 i,
                const T *__restrict valid_values,
                const u32 *__restrict offsets_bins,
                Tret *__restrict bin_averages) {
                u32 bin_start = offsets_bins[i];
                u32 bin_end   = offsets_bins[i + 1];
                u32 bin_count = bin_end - bin_start;
 
                auto for_each_values = [&](auto func) {
                    for (u32 j = bin_start; j < bin_end; j++) {
                        func(valid_values[j]);
                    }
                };
 
                bin_averages[i] = fct(for_each_values, bin_count);
            });
 
        return bin_compute;
    }
 
    template<class T>
    sham::DeviceBuffer<T> binned_sum(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges, // r bins
        u64 nbins,
        const sham::DeviceBuffer<T> &values, // ie f(r)
        const sham::DeviceBuffer<T> &keys,   // ie r
        u32 len) {
 
        return binned_compute<T, T>(
            sched, bin_edges, nbins, values, keys, len, [](auto for_each_values, u32 bin_count) {
                T sum = 0;
                for_each_values([&](T val) {
                    sum += val;
                });
                return sum;
            });
    }
 
    template<class T>
    sham::DeviceBuffer<T> binned_average(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges, // r bins
        u64 nbins,
        const sham::DeviceBuffer<T> &values, // ie f(r)
        const sham::DeviceBuffer<T> &keys,   // ie r
        u32 len) {
 
        return binned_compute<T, T>(
            sched,
            bin_edges,
            nbins,
            values,
            keys,
            len,
            [](auto for_each_values, u32 bin_count) -> T {
                T sum = T{};
                for_each_values([&](T val) {
                    sum += val;
                });
                if (bin_count == 0) {
                    return T{};
                } else {
                    return sum / bin_count;
                }
            });
    }
 
    template<class Tret, class T>
    sham::DeviceBuffer<Tret> device_histogram_mpi(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges,
        u64 nbins,
        const sham::DeviceBuffer<T> &values,
        u32 len) {
 
        auto local_counts = device_histogram<Tret, T>(sched, bin_edges, nbins, values, len);
 
        // local counts are now the global counts
        shamalgs::collective::reduce_buffer_in_place_sum(local_counts, MPI_COMM_WORLD);
 
        return local_counts;
    }
 
    template<class T>
    inline sham::DeviceBuffer<u64> device_histogram_u64_mpi(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges,
        u64 nbins,
        const sham::DeviceBuffer<T> &values,
        u32 len) {
        return device_histogram_mpi<u64, T>(sched, bin_edges, nbins, values, len);
    }
 
    template<class T>
    inline sham::DeviceBuffer<u32> device_histogram_u32_mpi(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges,
        u64 nbins,
        const sham::DeviceBuffer<T> &values,
        u32 len) {
        return device_histogram_mpi<u32, T>(sched, bin_edges, nbins, values, len);
    }
 
    template<class T>
    sham::DeviceBuffer<T> binned_sum_mpi(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges, // r bins
        u64 nbins,
        const sham::DeviceBuffer<T> &values, // ie f(r)
        const sham::DeviceBuffer<T> &keys,   // ie r
        u32 len) {
 
        auto local_result = binned_sum(sched, bin_edges, nbins, values, keys, len);
 
        // local result is now the global result
        shamalgs::collective::reduce_buffer_in_place_sum(local_result, MPI_COMM_WORLD);
 
        return local_result;
    }
 
    template<class T>
    sham::DeviceBuffer<T> binned_average_mpi(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges, // r bins
        u64 nbins,
        const sham::DeviceBuffer<T> &values, // ie f(r)
        const sham::DeviceBuffer<T> &keys,   // ie r
        u32 len,
        const sham::DeviceBuffer<u32> &bin_counts_global) {
 
        auto bin_sums = binned_sum_mpi(sched, bin_edges, nbins, values, keys, len);
 
        sham::kernel_call(
            shambase::get_check_ref(sched).get_queue(),
            sham::MultiRef{bin_counts_global},
            sham::MultiRef{bin_sums},
            nbins,
            [](u32 i, const u32 *__restrict bin_counts, T *__restrict bin_sums) {
                u32 bin_count = bin_counts[i];
                if (bin_count == 0) {
                    bin_sums[i] = T{};
                } else {
                    bin_sums[i] /= bin_count;
                }
            });
 
        return bin_sums;
    }
 
    template<class T>
    sham::DeviceBuffer<T> binned_average_mpi(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges, // r bins
        u64 nbins,
        const sham::DeviceBuffer<T> &values, // ie f(r)
        const sham::DeviceBuffer<T> &keys,   // ie r
        u32 len) {
 
        auto bin_counts = device_histogram_u32_mpi(sched, bin_edges, nbins, keys, len);
 
        // call the version with global bin counts pre-computed
        return binned_average_mpi(sched, bin_edges, nbins, values, keys, len, bin_counts);
    }
 
    template<class T>
    histogram_result<T> device_histogram_full_mpi(
        const sham::DeviceScheduler_ptr &sched,
        const sham::DeviceBuffer<T> &bin_edges,
        u64 nbins,
        const sham::DeviceBuffer<T> &values,
        u32 len) {
 
        SHAM_ASSERT(nbins > 1); // at least a sup and a inf
        SHAM_ASSERT(bin_edges.get_size() == nbins + 1);
 
        auto &q = shambase::get_check_ref(sched).get_queue();
 
        sham::DeviceBuffer<u64> counts
            = device_histogram_u64_mpi(sched, bin_edges, nbins, values, len);
 
        sham::DeviceBuffer<T> bins_center(nbins, sched);
        sham::DeviceBuffer<T> bins_width(nbins, sched);
 
        sham::kernel_call(
            q,
            sham::MultiRef{bin_edges},
            sham::MultiRef{bins_center, bins_width},
            nbins,
            [](u32 i,
               const T *__restrict bin_edges,
               T *__restrict bins_center,
               T *__restrict bins_width) {
                bins_center[i] = (bin_edges[i] + bin_edges[i + 1]) / 2;
                bins_width[i]  = bin_edges[i + 1] - bin_edges[i];
            });
 
        return {std::move(counts), std::move(bins_center), std::move(bins_width)};
    }
 
} // namespace shamalgs::numeric