pdi.hpp

// Copyright (C) The DDC development team, see COPYRIGHT.md file
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include <memory_resource>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <pdi.h>
 
#include "ddc/chunk_traits.hpp"
#include "ddc/discrete_vector.hpp"
 
namespace ddc {
 
template <class T>
static constexpr PDI_inout_t default_access_v
        = (std::is_lvalue_reference_v<T> && !std::is_const_v<std::remove_reference_t<T>>)
                  ? PDI_INOUT
                  : PDI_OUT;
 
template <class T>
static constexpr PDI_inout_t chunk_default_access_v = is_writable_chunk_v<T> ? PDI_INOUT : PDI_OUT;
 
class PdiEvent
{
    std::string m_event_name;
 
    std::vector<std::string> m_names;
 
    /// a memory buffer where temporary variables are stored until the class is destroyed
    std::pmr::monotonic_buffer_resource m_metadata;
 
public:
    explicit PdiEvent(std::string const& event_name) : m_event_name(event_name) {}
 
    PdiEvent(PdiEvent const& rhs) = delete;
 
    PdiEvent(PdiEvent&& rhs) noexcept = delete;
 
    ~PdiEvent() noexcept
    {
        PDI_event(m_event_name.c_str());
        for (std::string const& one_name : m_names) {
            PDI_reclaim(one_name.c_str());
        }
    }
 
    PdiEvent& operator=(PdiEvent const& rhs) = delete;
 
    PdiEvent& operator=(PdiEvent&& rhs) noexcept = delete;
 
    /// @{
    /// API with access argument
 
    template <
            PDI_inout_t access,
            class BorrowedChunk,
            std::enable_if_t<is_borrowed_chunk_v<BorrowedChunk>, int> = 0>
    PdiEvent& with(std::string const& name, BorrowedChunk&& data)
    {
        static_assert(
                !(access & PDI_IN) || (chunk_default_access_v<BorrowedChunk> & PDI_IN),
                "Invalid access for constant data");
        auto extents = detail::array(data.domain().extents());
        std::size_t& rank = *std::pmr::polymorphic_allocator<std::size_t>(&m_metadata).allocate(1);
        rank = extents.size();
        PDI_share((name + "_rank").c_str(), &rank, PDI_OUT);
        m_names.push_back(name + "_rank");
        PDI_share(
                (name + "_extents").c_str(),
                std::pmr::vector<std::size_t>(extents.begin(), extents.end(), &m_metadata).data(),
                PDI_OUT);
        m_names.push_back(name + "_extents");
        PDI_share(
                name.c_str(),
                const_cast<chunk_value_t<BorrowedChunk>*>(data.data_handle()),
                access);
        m_names.push_back(name);
        return *this;
    }
 
    template <
            PDI_inout_t access,
            class Arithmetic,
            std::enable_if_t<std::is_arithmetic_v<std::remove_reference_t<Arithmetic>>, int> = 0>
    PdiEvent& with(std::string const& name, Arithmetic&& data)
    {
        static_assert(
                !(access & PDI_IN) || (default_access_v<Arithmetic> & PDI_IN),
                "Invalid access for constant data");
        using value_type = std::remove_cv_t<std::remove_reference_t<Arithmetic>>;
        value_type* data_ptr = const_cast<value_type*>(&data);
        // for read-only data, we share a copy instead of the data itself in case we received a ref on a temporary,
        if constexpr (!(access & PDI_IN)) {
            data_ptr = std::pmr::polymorphic_allocator<value_type>(&m_metadata).allocate(1);
            *data_ptr = data;
        }
        PDI_share(name.c_str(), data_ptr, access);
        m_names.push_back(name);
        return *this;
    }
 
    template <PDI_inout_t access, class T>
    PdiEvent& and_with(std::string const& name, T&& t)
    {
        return with<access>(name, std::forward<T>(t));
    }
 
    /// @}
    /// API with access deduction
    /// @{
 
    /// Borrowed chunk overload (Chunk (const)& or ChunkSpan&& or ChunkSpan (const)&)
    template <class BorrowedChunk, std::enable_if_t<is_borrowed_chunk_v<BorrowedChunk>, int> = 0>
    PdiEvent& with(std::string const& name, BorrowedChunk&& data)
    {
        return with<chunk_default_access_v<BorrowedChunk>>(name, std::forward<BorrowedChunk>(data));
    }
 
    /// Arithmetic overload
    template <
            class Arithmetic,
            std::enable_if_t<std::is_arithmetic_v<std::remove_reference_t<Arithmetic>>, int> = 0>
    PdiEvent& with(std::string const& name, Arithmetic&& data)
    {
        return with<default_access_v<Arithmetic>>(name, std::forward<Arithmetic>(data));
    }
 
    /// With synonym
    template <class T>
    PdiEvent& and_with(std::string const& name, T&& t)
    {
        return with(name, std::forward<T>(t));
    }
 
    /// @}
};
 
template <PDI_inout_t access, class DataType>
void expose_to_pdi(std::string const& name, DataType&& data)
{
    PdiEvent(name).with<access>(name, std::forward<DataType>(data));
}
 
template <class DataType>
void expose_to_pdi(std::string const& name, DataType&& data)
{
    PdiEvent(name).with(name, std::forward<DataType>(data));
}
 
} // namespace ddc