window.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 "./datatypes.hpp"
#include "./group.hpp"
#include "./macros.hpp"
#include "./utils.hpp"
 
#include <mpi.h>
 
#include <algorithm>
#include <memory>
#include <tuple>
#include <utility>
 
namespace mpi {


  template <class BaseType> class window {
    public:
    using base_type = BaseType;

    window() = default;

    window(window const &) = delete;

    window &operator=(window const &) = delete;

    window(window &&other) noexcept
       : win_{std::exchange(other.win_, MPI_WIN_NULL)},
         comm_{std::exchange(other.comm_, communicator{MPI_COMM_NULL})},
         owned_{std::exchange(other.owned_, false)},
         data_{std::exchange(other.data_, nullptr)},
         size_{std::exchange(other.size_, 0)} {}

    window &operator=(window &&rhs) noexcept {
      if (this != std::addressof(rhs)) {
        free();
        win_   = std::exchange(rhs.win_, MPI_WIN_NULL);
        comm_  = std::exchange(rhs.comm_, communicator{MPI_COMM_NULL});
        owned_ = std::exchange(rhs.owned_, false);
        data_  = std::exchange(rhs.data_, nullptr);
        size_  = std::exchange(rhs.size_, 0);
      }
      return *this;
    }

    explicit window(communicator const &c, BaseType *base_ptr, MPI_Aint sz, MPI_Info info = MPI_INFO_NULL)
       : comm_(c.get()), data_(base_ptr), size_(sz) {
      ASSERT(size_ >= 0)
      ASSERT(!(data_ == nullptr && size_ > 0))
      if (has_env) check_mpi_call(MPI_Win_create(data_, size_ * sizeof(BaseType), sizeof(BaseType), info, c.get(), &win_), "MPI_Win_create");
    }

    explicit window(communicator const &c, MPI_Aint sz, MPI_Info info = MPI_INFO_NULL) : comm_(c.get()), size_(sz) {
      ASSERT(size_ >= 0)
      if (has_env) {
        check_mpi_call(MPI_Win_allocate(size_ * sizeof(BaseType), sizeof(BaseType), info, c.get(), &data_, &win_), "MPI_Win_allocate");
      } else {
        data_ = new BaseType[size_]; // NOLINT (new is fine here)
      }
      owned_ = true;
    }

    explicit operator MPI_Win() const { return win_; };

    explicit operator MPI_Win *() { return &win_; };

    virtual ~window() { free(); }

    void free() noexcept {
      if (has_env) {
        if (win_ != MPI_WIN_NULL) {
          MPI_Win_fence(0, win_);
          MPI_Win_free(&win_);
        }
      } else if (owned_) {
        delete[] data_;
      }
      owned_ = false;
      data_  = nullptr;
      size_  = 0;
    }

    void fence(int assert = 0) const {
      if (has_env) check_mpi_call(MPI_Win_fence(assert, win_), "MPI_Win_fence");
    }

    void flush(int rank = -1) const {
      if (has_env) {
        if (rank < 0) {
          check_mpi_call(MPI_Win_flush_all(win_), "MPI_Win_flush_all");
        } else {
          check_mpi_call(MPI_Win_flush(rank, win_), "MPI_Win_flush");
        }
      }
    }

    void sync() const {
      if (has_env) check_mpi_call(MPI_Win_sync(win_), "MPI_Win_sync");
    }

    void lock(int rank = -1, int lock_type = MPI_LOCK_SHARED, int assert = 0) const {
      if (has_env) {
        if (rank < 0) {
          check_mpi_call(MPI_Win_lock_all(assert, win_), "MPI_Win_lock_all");
        } else {
          check_mpi_call(MPI_Win_lock(lock_type, rank, assert, win_), "MPI_Win_lock");
        }
      }
    }

    void unlock(int rank = -1) const {
      if (has_env) {
        if (rank < 0) {
          check_mpi_call(MPI_Win_unlock_all(win_), "MPI_Win_unlock_all");
        } else {
          check_mpi_call(MPI_Win_unlock(rank, win_), "MPI_Win_unlock");
        }
      }
    }

    void start(group const &grp, int assert = 0) const {
      if (has_env) check_mpi_call(MPI_Win_start(grp.get(), assert, win_), "MPI_Win_start");
    }

    void complete() const {
      if (has_env) check_mpi_call(MPI_Win_complete(win_), "MPI_Win_complete");
    }

    void post(group const &grp, int assert = 0) const {
      if (has_env) check_mpi_call(MPI_Win_post(grp.get(), assert, win_), "MPI_Win_post");
    }

    void wait() const {
      if (has_env) check_mpi_call(MPI_Win_wait(win_), "MPI_Win_wait");
    }

    template <typename TargetType = BaseType, typename OriginType>
      requires(has_mpi_type<OriginType> && has_mpi_type<TargetType>)
    void get(OriginType *origin_addr, int origin_count, int target_rank, MPI_Aint target_disp = 0, int target_count = -1) const {
      ASSERT(origin_count >= 0 && target_disp >= 0);
      target_count = target_count < 0 ? origin_count : target_count;
      if (has_env) {
        auto origin_datatype = mpi_type<OriginType>::get();
        auto target_datatype = mpi_type<TargetType>::get();
        check_mpi_call(MPI_Get(origin_addr, origin_count, origin_datatype, target_rank, target_disp, target_count, target_datatype, win_), "MPI_Get");
      } else {
        ASSERT(target_rank == 0); // there are no other ranks
        std::copy(data_, data_ + target_count, origin_addr);
      }
    }

    template <typename TargetType = BaseType, typename OriginType>
      requires(has_mpi_type<OriginType> && has_mpi_type<TargetType>)
    void put(OriginType *origin_addr, int origin_count, int target_rank, MPI_Aint target_disp = 0, int target_count = -1) const {
      ASSERT(origin_count >= 0 && target_disp >= 0);
      target_count = target_count < 0 ? origin_count : target_count;
      if (has_env) {
        auto origin_datatype = mpi_type<OriginType>::get();
        auto target_datatype = mpi_type<TargetType>::get();
        check_mpi_call(MPI_Put(origin_addr, origin_count, origin_datatype, target_rank, target_disp, target_count, target_datatype, win_), "MPI_Put");
      } else {
        ASSERT(target_rank == 0); // there are no other ranks
        std::copy(origin_addr, origin_addr + origin_count, data_);
      }
    }

    [[nodiscard]] BaseType *base() const { return data_; }

    [[nodiscard]] MPI_Aint size() const { return size_; }

    [[nodiscard]] int disp_unit() const { return sizeof(BaseType); }

    [[nodiscard]] communicator get_communicator() const { return comm_; }
 
    protected:
    MPI_Win win_{MPI_WIN_NULL};
    communicator comm_{MPI_COMM_NULL};
    bool owned_{false};
    BaseType *data_{nullptr};
    MPI_Aint size_{0};
  };

  template <class BaseType> class shared_window : public window<BaseType> {
    public:
    shared_window() = default;

    explicit shared_window(shared_communicator const &c, MPI_Aint sz, MPI_Info info = MPI_INFO_NULL) {
      ASSERT(sz >= 0)
      comm_ = c.get();
      size_ = sz;
      if (has_env) {
        check_mpi_call(MPI_Win_allocate_shared(size_ * sizeof(BaseType), sizeof(BaseType), info, c.get(), &data_, &win_), "MPI_Win_allocate_shared");
      } else {
        data_ = new BaseType[size_]; // NOLINT (new is fine here)
      }
      owned_ = true;
    }

    [[nodiscard]] std::tuple<MPI_Aint, int, void *> query(int rank = MPI_PROC_NULL) const {
      if (has_env) {
        MPI_Aint sz   = 0;
        int du        = 0;
        void *baseptr = nullptr;
        check_mpi_call(MPI_Win_shared_query(win_, rank, &sz, &du, &baseptr), "MPI_Win_shared_query");
        return {sz, du, baseptr};
      } else {
        return {size_, sizeof(BaseType), data_};
      }
    }

    [[nodiscard]] BaseType *base(int rank = MPI_PROC_NULL) const { return static_cast<BaseType *>(std::get<2>(query(rank))); }

    [[nodiscard]] MPI_Aint size(int rank = MPI_PROC_NULL) const { return std::get<0>(query(rank)) / sizeof(BaseType); }

    [[nodiscard]] int disp_unit(int rank = MPI_PROC_NULL) const { return std::get<1>(query(rank)); }

    [[nodiscard]] shared_communicator get_communicator() const { return comm_.get(); }
 
    private:
    using window<BaseType>::win_;
    using window<BaseType>::comm_;
    using window<BaseType>::owned_;
    using window<BaseType>::data_;
    using window<BaseType>::size_;
  };

 
} // namespace mpi