broadcast.hpp

Go to the documentation of this file.

// Copyright (c) 2024--present, The Simons Foundation
// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
// SPDX-License-Identifier: Apache-2.0
// See LICENSE in the root of this distribution for details.

 
#pragma once
 
#include "../basic_functions.hpp"
#include "../traits.hpp"
 
#include <mpi/mpi.hpp>
 
#include <cstddef>
#include <span>
 
namespace nda {

  template <typename A>
  void mpi_broadcast(A &&a, mpi::communicator comm = {}, int root = 0) // NOLINT (temporary views are allowed here)
    requires(is_regular_or_view_v<A>)
  {
    // broadcast the shape of the input array/view on the root process
    auto dims = a.shape();
    mpi::broadcast(dims, comm, root);
 
    // resize or check the output array/view on non-root processes
    if (comm.rank() != root) resize_or_check_if_view(a, dims);
 
    // call mpi::broadcast_range with a span if the array/view is contiguous with positive strides
    if (a.is_contiguous() and a.has_positive_strides()) {
      auto a_span = std::span{a.data(), static_cast<std::size_t>(a.size())};
      mpi::broadcast_range(a_span, comm, root);
    } else {
      mpi::broadcast_range(a, comm, root);
    }
  }
 
} // namespace nda