traits.hpp

/*******************************************************************************
 * Copyright (C) 2025-2026 Theodore Chang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 ******************************************************************************/
 
#ifndef TRAITS_HPP
#define TRAITS_HPP
 
#include "ezp.h"
 
#include <array>
#include <atomic>
#include <cmath>
#include <ranges>
#include <vector>
 
namespace ezp {
    template<typename T> constexpr auto always_false_v = false;
 
    template<typename T> concept floating_t = std::is_same_v<T, float> || std::is_same_v<T, double>;
    template<typename T> concept complex_t = std::is_same_v<T, std::complex<typename T::value_type>> && floating_t<typename T::value_type>;
    template<typename T> concept data_t = floating_t<T> || complex_t<T>;
    template<typename T> concept index_t = std::is_same_v<T, std::int32_t> || std::is_same_v<T, std::int64_t>;
 
    template<typename T> concept has_mem = requires(T t) { requires data_t<std::remove_pointer_t<decltype(t.mem())>>; };
    template<typename T> concept has_memptr = requires(T t) { requires data_t<std::remove_pointer_t<decltype(t.memptr())>>; };
    template<typename T> concept has_data_method = requires(T t) { requires data_t<std::remove_pointer_t<decltype(t.data())>>; };
    template<typename T> concept has_data_member = requires(T t) { requires data_t<std::remove_pointer_t<decltype(t.data)>>; };
    template<typename T> concept has_iterator = requires(T t) {
        requires std::ranges::contiguous_range<T>;
        requires data_t<std::remove_reference_t<decltype(*t.begin())>>;
    };
    template<typename T> concept has_data_pointer = has_mem<T> || has_memptr<T> || has_data_method<T> || has_iterator<T>;
 
    template<typename T> concept full_container_t = requires(T t) { requires index_t<decltype(t.n_rows)> && index_t<decltype(t.n_cols)> && has_data_pointer<T>; };
    template<typename T> concept band_container_t = requires(T t) { requires index_t<decltype(t.kl)> && index_t<decltype(t.ku)> && full_container_t<T>; };
    template<typename T> concept band_symm_container_t = requires(T t) { requires index_t<decltype(t.klu)> && full_container_t<T>; };
 
    template<typename T> struct WorkType {
        static_assert(!std::is_same_v<T, T>, "not defined");
    };
    template<> struct WorkType<double> {
        using type = double;
    };
    template<> struct WorkType<float> {
        using type = float;
    };
    template<> struct WorkType<complex8> {
        using type = float;
    };
    template<> struct WorkType<complex16> {
        using type = double;
    };
    template<typename T> using work_t = typename WorkType<T>::type;
 
    template<data_t DT, index_t IT> struct base_mat {
        using data_type = DT;
        using index_type = IT;
 
        IT n_rows, n_cols;
        DT* data;
        bool distributed{};
 
        base_mat() = default;
 
        base_mat(const IT rows, const IT cols, DT* const ptr, const bool dist = false)
            : n_rows(rows)
            , n_cols(cols)
            , data(ptr)
            , distributed(dist) {}
    };
 
    template<data_t DT, index_t IT> struct full_mat : base_mat<DT, IT> {
        using base_mat<DT, IT>::base_mat;
    };
 
    template<data_t DT, index_t IT> struct band_mat : base_mat<DT, IT> {
        IT kl, ku;
 
        band_mat() = default;
 
        band_mat(const IT rows, const IT cols, const IT kl, const IT ku, DT* const ptr, const bool dist = false)
            : base_mat<DT, IT>(rows, cols, ptr, dist)
            , kl(kl)
            , ku(ku) {}
    };
 
    template<data_t DT, index_t IT> struct band_symm_mat : base_mat<DT, IT> {
        IT klu;
 
        band_symm_mat() = default;
 
        band_symm_mat(const IT rows, const IT cols, const IT klu, DT* const ptr, const bool dist = false)
            : base_mat<DT, IT>(rows, cols, ptr, dist)
            , klu(klu) {}
    };
 
    template<typename T> concept wrapper_t = requires(T t) {
        requires std::is_same_v<T, full_mat<typename T::data_type, typename T::index_type>> || std::is_same_v<T, band_mat<typename T::data_type, typename T::index_type>> || std::is_same_v<T, band_symm_mat<typename T::data_type, typename T::index_type>>;
    };
 
    template<index_t IT> using desc = std::array<IT, 9>;
 
    class blacs_base {
#ifdef EZP_RELEASE_ONCE
    protected:
        static std::atomic_bool RELEASED;
#endif
    };
 
#ifdef EZP_RELEASE_ONCE
#define EZP_ENSURE_SAFE_EXIT std::atomic_bool ezp::blacs_base::RELEASED = false;
#endif
 
    template<index_t IT = int_t> class blacs_env final : blacs_base {
        static constexpr IT ZERO{0}, ONE{1};
 
        static std::atomic_bool FINALIZE;
 
        IT _rank{-1}, _size{-1};
 
    public:
        blacs_env() { blacs_pinfo(&_rank, &_size); }
 
        ~blacs_env() {
#ifdef EZP_RELEASE_ONCE
            if(RELEASED.exchange(true)) return;
#endif
            blacs_exit(FINALIZE ? &ZERO : &ONE);
        }
 
        static void do_not_manage_mpi() { FINALIZE = false; }
 
        auto rank() const { return _rank; }
        auto size() const { return _size; }
    };
 
    template<index_t IT> std::atomic_bool blacs_env<IT>::FINALIZE{true};
 
    template<index_t IT = int_t> const auto& get_env() {
        static const blacs_env<IT> scoped_env;
 
        return scoped_env;
    }
 
    template<index_t IT> class blacs_context final {
        static constexpr IT ZERO{0}, ONE{1}, NEGONE{-1};
        static constexpr char SCOPE = 'A', TOP = ' ';
 
        char layout{'R'};
 
        IT info{-1};
 
        auto init() {
            blacs_get(&NEGONE, &ZERO, &context);
            blacs_gridinit(&context, &layout, &n_rows, &n_cols);
            blacs_pinfo(&rank, &size);
            blacs_gridinfo(&context, &n_rows, &n_cols, &my_row, &my_col);
        }
 
        auto release() {
            if(context >= 0) blacs_gridexit(&context);
        }
 
        template<data_t DT> auto copy_to(const DT* A, const IT* desc_a, DT* B, const IT* desc_b) {
            // ReSharper disable CppCStyleCast
            if(std::is_same_v<DT, double>) {
                using E = double;
                pdgemr2d(desc_a + 2, desc_a + 3, (E*)A, &ONE, &ONE, desc_a, (E*)B, &ONE, &ONE, desc_b, &context);
            }
            else if(std::is_same_v<DT, float>) {
                using E = float;
                psgemr2d(desc_a + 2, desc_a + 3, (E*)A, &ONE, &ONE, desc_a, (E*)B, &ONE, &ONE, desc_b, &context);
            }
            else if(std::is_same_v<DT, complex16>) {
                using E = complex16;
                pzgemr2d(desc_a + 2, desc_a + 3, (E*)A, &ONE, &ONE, desc_a, (E*)B, &ONE, &ONE, desc_b, &context);
            }
            else if(std::is_same_v<DT, complex8>) {
                using E = complex8;
                pcgemr2d(desc_a + 2, desc_a + 3, (E*)A, &ONE, &ONE, desc_a, (E*)B, &ONE, &ONE, desc_b, &context);
            }
            // ReSharper restore CppCStyleCast
        }
 
    public:
        IT n_rows, n_cols, context{-1}, rank{-1}, size{-1}, my_row{-1}, my_col{-1};
 
        blacs_context()
            : blacs_context(get_env<IT>().size(), 1) {}
 
        explicit blacs_context(const char order)
            : layout(order)
            , n_rows(-1)
            , n_cols(-1) {
            const auto& env = get_env<IT>();
            n_rows = std::max(IT{1}, static_cast<IT>(std::sqrt(env.size())));
            n_cols = env.size() / n_rows;
            init();
        }
 
        blacs_context(const IT rows, const IT cols, const char order = 'R')
            : layout(order)
            , n_rows(rows)
            , n_cols(cols) { init(); }
 
        blacs_context(const blacs_context& other)
            : layout(other.layout)
            , n_rows(other.n_rows)
            , n_cols(other.n_cols) { init(); }
 
        blacs_context(blacs_context&&) noexcept = delete;
        blacs_context& operator=(const blacs_context&) = delete;
        blacs_context& operator=(blacs_context&&) noexcept = delete;
 
        ~blacs_context() { release(); }
 
        auto desc_g(const IT num_rows, const IT num_cols) {
            desc<IT> desc_t{};
 
            descinit(desc_t.data(), &num_rows, &num_cols, &num_rows, &num_cols, &ZERO, &ZERO, &context, &num_rows, &info);
 
            return desc_t;
        }
 
        auto desc_l(const IT num_rows, const IT num_cols, const IT row_block, const IT col_block, const IT lead) {
            desc<IT> desc_t{};
 
            const auto loc_lead = std::max(IT{1}, lead);
            descinit(desc_t.data(), &num_rows, &num_cols, &row_block, &col_block, &ZERO, &ZERO, &context, &loc_lead, &info);
 
            return desc_t;
        }
 
        auto desc_l(const IT num_rows, const IT num_cols, const IT block, const IT lead) { return desc_l(num_rows, num_cols, block, block, lead); }
 
        template<data_t DT> auto scatter(const full_mat<DT, IT>& A, const desc<IT>& desc_a, std::vector<DT>& B, const desc<IT>& desc_b) {
            if(!A.distributed) return copy_to(A.data, desc_a.data(), B.data(), desc_b.data());
 
            for(auto i = 0u; i < B.size(); ++i) B[i] = A.data[i];
        }
 
        template<data_t DT> auto gather(const std::vector<DT>& A, const desc<IT>& desc_a, const full_mat<DT, IT>& B, const desc<IT>& desc_b) {
            if(!B.distributed) return copy_to(A.data(), desc_a.data(), B.data, desc_b.data());
 
            for(auto i = 0u; i < A.size(); ++i) B.data[i] = A[i];
        }
 
        auto copy_to(const IT* A, const IT* desc_a, IT* B, const IT* desc_b) { pigemr2d(desc_a + 2, desc_a + 3, A, &ONE, &ONE, desc_a, B, &ONE, &ONE, desc_b, &context); }
 
        [[nodiscard]] bool is_valid() const { return my_row >= 0 && my_col >= 0; }
 
        auto row_block(const IT n) const { return std::max(IT{1}, n / n_rows); }
 
        auto col_block(const IT n) const { return std::max(IT{1}, n / n_cols); }
 
        auto rows(const IT n, const IT nb) const { return numroc(&n, &nb, &my_row, &ZERO, &n_rows); }
 
        auto cols(const IT n, const IT nb) const { return numroc(&n, &nb, &my_col, &ZERO, &n_cols); }
 
        IT amx(IT number) const {
            igamx2d(&context, &SCOPE, &TOP, &ONE, &ONE, &number, &ONE, nullptr, nullptr, &NEGONE, &NEGONE, &NEGONE);
            return number;
        }
 
        IT amn(IT number) const {
            igamn2d(&context, &SCOPE, &TOP, &ONE, &ONE, &number, &ONE, nullptr, nullptr, &NEGONE, &NEGONE, &NEGONE);
            return number;
        }
    };
} // namespace ezp
 
#endif // TRAITS_HPP