datatypes.hpp

Go to the documentation of this file.

// Copyright (c) 2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Alexander Hampel, Olivier Parcollet, Nils Wentzell

 
#pragma once
 
#include "./utils.hpp"
 
#include <mpi.h>
 
#include <algorithm>
#include <array>
#include <complex>
#include <cstdlib>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace mpi {


  template <typename T> struct mpi_type {};
 
#define D(T, MPI_TY)                                                                                                                                 \
                                                                                                    \
  template <> struct mpi_type<T> {                                                                                                                   \
    [[nodiscard]] static MPI_Datatype get() noexcept { return MPI_TY; }                                                                              \
  }
 
  // mpi_type specialization for various built-in types
  D(bool, MPI_C_BOOL);
  D(char, MPI_CHAR);
  D(int, MPI_INT);
  D(long, MPI_LONG);
  D(long long, MPI_LONG_LONG);
  D(double, MPI_DOUBLE);
  D(float, MPI_FLOAT);
  D(std::complex<double>, MPI_C_DOUBLE_COMPLEX);
  D(unsigned long, MPI_UNSIGNED_LONG);
  D(unsigned int, MPI_UNSIGNED);
  D(unsigned long long, MPI_UNSIGNED_LONG_LONG);
#undef D

  template <typename E>
    requires(std::is_enum_v<E>)
  struct mpi_type<E> : mpi_type<std::underlying_type_t<E>> {};

  template <typename T> struct mpi_type<const T> : mpi_type<T> {};

  template <typename T, typename = void> constexpr bool has_mpi_type = false;

  template <typename T> constexpr bool has_mpi_type<T, std::void_t<decltype(mpi_type<T>::get())>> = true;
 
  namespace detail {
 
    // Helper struct to check if member types are mpi-serializable, i.e. have an associated mpi_type
    struct serialize_checker {
      template <typename T>
      void operator&(T &)
        requires(has_mpi_type<T>)
      {}
    };
 
  } // namespace detail

  template <typename T>
  concept Serializable = requires(const T ac, T a, detail::serialize_checker ar) {
    { ac.serialize(ar) } -> std::same_as<void>;
    { a.deserialize(ar) } -> std::same_as<void>;
  };

  template <typename... Ts> [[nodiscard]] MPI_Datatype get_mpi_type(std::tuple<Ts...> tup) {
    static constexpr int N            = sizeof...(Ts);
    std::array<MPI_Datatype, N> types = {mpi_type<std::remove_reference_t<Ts>>::get()...};
 
    // the number of elements per type (we want 1 per type)
    std::array<int, N> blocklen;
    for (int i = 0; i < N; ++i) { blocklen[i] = 1; }
 
    // displacements of the blocks in bytes w.r.t. to the memory address of the first block
    std::array<MPI_Aint, N> disp;
    // initialize displacement array from the tuple element addresses
    []<size_t... Is>(std::index_sequence<Is...>, auto &t, MPI_Aint *d) {
      ((d[Is] = (char *)&std::get<Is>(t) - (char *)&std::get<0>(t)), ...);
      // account for non-trivial memory layouts of the tuple elements
      auto min_el = *std::min_element(d, d + sizeof...(Ts));
      ((d[Is] -= min_el), ...);
    }(std::index_sequence_for<Ts...>{}, tup, disp.data());
 
    // create and return MPI datatype
    MPI_Datatype cty{};
    check_mpi_call(MPI_Type_create_struct(N, blocklen.data(), disp.data(), types.data(), &cty), "MPI_Type_create_struct");
    check_mpi_call(MPI_Type_commit(&cty), "MPI_Type_commit");
    return cty;
  }

  template <typename... Ts> struct mpi_type<std::tuple<Ts...>> {
    [[nodiscard]] static MPI_Datatype get() noexcept {
      static MPI_Datatype type = get_mpi_type(std::tuple<Ts...>{});
      return type;
    }
  };

  template <typename U>
    requires(not Serializable<U>) and requires(U u) { tie_data(u); }
  struct mpi_type<U> {
    [[nodiscard]] static MPI_Datatype get() noexcept {
      static MPI_Datatype type = get_mpi_type(tie_data(U{}));
      return type;
    }
  };
 
  namespace detail {
 
    // Archive helper class to obtain MPI custom type info using references to class members.
    struct mpi_archive {
      std::vector<int> block_lengths{};
      std::vector<MPI_Aint> displacements{};
      std::vector<MPI_Datatype> types{};
      MPI_Aint base_address{};
 
      // Constructor sets the base address of the object.
      explicit mpi_archive(const void *base) { MPI_Get_address(base, &base_address); }
 
      // Overloaded operator& to process members to set the block lengths, displacements and MPI types.
      template <typename T>
        requires(has_mpi_type<T>)
      mpi_archive &operator&(const T &member) {
        types.push_back(mpi_type<T>::get());
        MPI_Aint address{};
        MPI_Get_address(&member, &address);
        displacements.push_back(MPI_Aint_diff(address, base_address));
        block_lengths.push_back(1);
        return *this;
      }
    };
 
  } // namespace detail

  template <Serializable T> [[nodiscard]] MPI_Datatype get_mpi_type(const T &obj) {
    detail::mpi_archive ar(&obj);
    obj.serialize(ar);
    MPI_Datatype mpi_type{};
    MPI_Type_create_struct(static_cast<int>(ar.block_lengths.size()), ar.block_lengths.data(), ar.displacements.data(), ar.types.data(), &mpi_type);
    MPI_Type_commit(&mpi_type);
    return mpi_type;
  }

  template <Serializable S> struct mpi_type<S> {
    [[nodiscard]] static MPI_Datatype get() noexcept {
      static MPI_Datatype type = get_mpi_type(S{});
      return type;
    }
  };

 
} // namespace mpi