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 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;
 
  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... T> struct mpi_type<std::tuple<T...>> {
    [[nodiscard]] static MPI_Datatype get() noexcept { return get_mpi_type(std::tuple<T...>{}); }
  };
 
  template <typename T> struct mpi_type_from_tie {
    [[nodiscard]] static MPI_Datatype get() noexcept { return get_mpi_type(tie_data(T{})); }
  };
 
} // namespace mpi