chunk.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 "./communicator.hpp"
 
#include <itertools/itertools.hpp>
 
#include <iterator>
#include <stdexcept>
#include <utility>
 
namespace mpi {
 
  [[nodiscard]] inline long chunk_length(long end, int nranges, int i, long min_size = 1) {
    if (min_size < 1 || end % min_size != 0) throw std::runtime_error("Error in mpi::chunk_length: min_size must be a divisor of end");
    auto [node_begin, node_end] = itertools::chunk_range(0, end / min_size, nranges, i);
    return (node_end - node_begin) * min_size;
  }
 
  template <typename R> [[nodiscard]] auto chunk(R &&rg, communicator c = {}) {
    auto total_size           = itertools::distance(std::cbegin(rg), std::cend(rg));
    auto [start_idx, end_idx] = itertools::chunk_range(0, total_size, c.size(), c.rank());
    return itertools::slice(std::forward<R>(rg), start_idx, end_idx);
  }
 
} // namespace mpi