boundaryData.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2024 - 2026 NeoN authors
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include "NeoN/core/primitives/label.hpp"
#include "NeoN/core/primitives/scalar.hpp"
#include "NeoN/core/executor/executor.hpp"
#include "NeoN/core/vector/vector.hpp"
 
#include <vector>
#include <utility>
#include <cstdio>
#include <algorithm>
 
#ifdef NF_WITH_MPI_SUPPORT
#include <mpi.h>
#include <optional>
#include "NeoN/core/mpi/environment.hpp"
#include "NeoN/core/mpi/operators.hpp"
#include "NeoN/core/parallelAlgorithms.hpp"
#endif
 
namespace NeoN
{
 
 
template<typename ValueType>
class BoundaryData
{
 
public:
 
    using BoundaryDataType = ValueType;
 
    ~BoundaryData()
    {
        // commBuffers_ backs the memory of any in-flight MPI_Isend/MPI_Irecv calls.
        // Destroying them while operations are pending is undefined behaviour, so
        // drain all outstanding requests before the storage is freed.
        waitAll();
    }
 
    BoundaryData(const BoundaryData<ValueType>& rhs)
        : exec_(rhs.exec_), value_(rhs.value_), refValue_(rhs.refValue_),
          valueFraction_(rhs.valueFraction_), refGrad_(rhs.refGrad_),
          boundaryTypes_(rhs.boundaryTypes_), offset_(rhs.offset_), nBoundaries_(rhs.nBoundaries_),
          nBoundaryFaces_(rhs.nBoundaryFaces_)
    {}
 
 
    BoundaryData(const Executor& exec, const BoundaryData<ValueType>& rhs)
        : exec_(rhs.exec_), value_(exec, rhs.value_), refValue_(exec, rhs.refValue_),
          valueFraction_(exec, rhs.valueFraction_), refGrad_(exec, rhs.refGrad_),
          boundaryTypes_(exec, rhs.boundaryTypes_), offset_(SerialExecutor {}, rhs.offset_),
          nBoundaries_(rhs.nBoundaries_), nBoundaryFaces_(rhs.nBoundaryFaces_)
    {}
 
 
    BoundaryData(const Executor& exec, localIdx nBoundaryFaces, localIdx nBoundaryTypes)
        : exec_(exec), value_(exec, nBoundaryFaces, ValueType {}),
          refValue_(exec, nBoundaryFaces, ValueType {}),
          valueFraction_(exec, nBoundaryFaces, scalar(0)),
          refGrad_(exec, nBoundaryFaces, ValueType {}), boundaryTypes_(exec, nBoundaryTypes),
          offset_(SerialExecutor {}, nBoundaryTypes + 1), nBoundaries_(nBoundaryTypes),
          nBoundaryFaces_(nBoundaryFaces)
    {}
 
    BoundaryData(const Executor& exec, const std::vector<localIdx>& offsets)
        : BoundaryData(exec, offsets.back(), static_cast<localIdx>(offsets.size() - 1))
    {
        offset_ = Vector(SerialExecutor {}, offsets);
    }
 
 
    const Vector<ValueType>& value() const
    {
        waitAll();
        return value_;
    }
 
    Vector<ValueType>& value()
    {
        waitAll();
        return value_;
    }
 
    const Vector<ValueType>& refValue() const { return refValue_; }
 
    Vector<ValueType>& refValue() { return refValue_; }
 
    const Vector<scalar>& valueFraction() const { return valueFraction_; }
 
    Vector<scalar>& valueFraction() { return valueFraction_; }
 
    const Vector<ValueType>& refGrad() const { return refGrad_; }
 
    Vector<ValueType>& refGrad() { return refGrad_; }
 
    const Vector<int>& boundaryTypes() const { return boundaryTypes_; }
 
    const Vector<localIdx>& offset() const { return offset_; }
 
    localIdx nBoundaries() const { return nBoundaries_; }
 
    localIdx nBoundaryFaces() const { return nBoundaryFaces_; }
 
    localIdx nBoundaryFaces(localIdx patchId) const
    {
        return offset_.data()[patchId + 1] - offset_.data()[patchId];
    }
 
    const Executor& exec() { return exec_; }
 
    BoundaryData<ValueType>& operator=(const BoundaryData<ValueType>& rhs)
    {
 
        // TODO maybe dont overwrite nBoundaries and nBoundaryFaces
        // but use them for a sanity check
        nBoundaries_ = rhs.nBoundaries_;
        nBoundaryFaces_ = rhs.nBoundaryFaces_;
 
        value_ = rhs.value_;
        refValue_ = rhs.refValue_;
        valueFraction_ = rhs.valueFraction_;
        refGrad_ = rhs.refGrad_;
        boundaryTypes_ = rhs.boundaryTypes_;
        offset_ = rhs.offset_;
        return *this;
    }
 
    BoundaryData<ValueType>& operator=(const BoundaryData<ValueType>&& rhs)
    {
 
        // TODO maybe dont overwrite nBoundaries and nBoundaryFaces
        // but use them for a sanity check
        nBoundaries_ = rhs.nBoundaries_;
        nBoundaryFaces_ = rhs.nBoundaryFaces_;
 
        value_ = std::move(rhs.value_);
        refValue_ = std::move(rhs.refValue_);
        valueFraction_ = std::move(rhs.valueFraction_);
        refGrad_ = std::move(rhs.refGrad_);
        boundaryTypes_ = std::move(rhs.boundaryTypes_);
        offset_ = std::move(rhs.offset_);
        return *this;
    }
 
#ifdef NF_WITH_MPI_SUPPORT
    void communicate(std::pair<localIdx, localIdx> range, int neighborRank)
    {
        const auto [rangeStart, rangeEnd] = range;
        const localIdx patchSize = rangeEnd - rangeStart;
 
        mpi::Environment mpiEnv;
        CommBuffer buf;
        buf.rangeStart = rangeStart;
        buf.patchSize = patchSize;
 
        const auto byteCount =
            static_cast<mpi_label_t>(patchSize) * static_cast<mpi_label_t>(sizeof(ValueType));
        const auto neighborRankLabel = static_cast<mpi_label_t>(neighborRank);
 
        // Deterministic, symmetric tag for the processor patch shared by (myRank, neighborRank).
        // A unique tag per unordered rank pair, identical on both sides, makes each isend/irecv
        // match its true partner regardless of posting order. min*P+max is symmetric so both
        // ranks of the pair compute the same tag.
        const auto nProcs = static_cast<mpi_label_t>(mpiEnv.sizeRank());
        const auto myRankLabel = static_cast<mpi_label_t>(mpiEnv.rank());
        const mpi_label_t pairTag = std::min(myRankLabel, neighborRankLabel) * nProcs
                                  + std::max(myRankLabel, neighborRankLabel);
 
        const bool useGpuPath = mpiEnv.gpuAwareMpi() && std::holds_alternative<GPUExecutor>(exec_);
 
        MPI_Request sendReq, recvReq;
        if (useGpuPath)
        {
            buf.deviceRecvBuf = Vector<ValueType>(exec_, patchSize, ValueType {});
            mpi::isend<char>(
                reinterpret_cast<const char*>(value_.data() + rangeStart),
                byteCount,
                neighborRankLabel,
                pairTag,
                mpiEnv.comm(),
                &sendReq
            );
            mpi::irecv<char>(
                reinterpret_cast<char*>(buf.deviceRecvBuf->data()),
                byteCount,
                neighborRankLabel,
                pairTag,
                mpiEnv.comm(),
                &recvReq
            );
        }
        else
        {
            auto valH = value_.copyToHost();
            buf.sendBuf.resize(static_cast<std::size_t>(patchSize));
            buf.recvBuf.resize(static_cast<std::size_t>(patchSize));
            for (localIdx k = 0; k < patchSize; k++)
                buf.sendBuf[static_cast<std::size_t>(k)] = valH.view()[rangeStart + k];
            mpi::isend<char>(
                reinterpret_cast<const char*>(buf.sendBuf.data()),
                byteCount,
                neighborRankLabel,
                pairTag,
                mpiEnv.comm(),
                &sendReq
            );
            mpi::irecv<char>(
                reinterpret_cast<char*>(buf.recvBuf.data()),
                byteCount,
                neighborRankLabel,
                pairTag,
                mpiEnv.comm(),
                &recvReq
            );
        }
        communicating_ = true;
        requests_.push_back(sendReq);
        requests_.push_back(recvReq);
        commBuffers_.push_back(std::move(buf));
    }
 
    bool isComplete() const
    {
        if (requests_.empty() || !communicating_) return true;
        for (auto& req : requests_)
        {
            if (!mpi::test(&req)) return false;
        }
        communicating_ = false;
        return true;
    }
 
 
#endif
 
    void waitAll() const
    {
#ifdef NF_WITH_MPI_SUPPORT
        if (requests_.empty() || !communicating_) return;
        while (!isComplete())
        {
        }
        mpi::Environment mpiEnv;
        const bool useGpuPath = mpiEnv.gpuAwareMpi() && std::holds_alternative<GPUExecutor>(exec_);
        if (useGpuPath)
        {
            for (const auto& buf : commBuffers_)
            {
                auto srcView = buf.deviceRecvBuf->view();
                auto dstView = value_.view();
                const localIdx start = buf.rangeStart;
                parallelFor(
                    exec_,
                    {0, buf.patchSize},
                    KOKKOS_LAMBDA(const localIdx k) { dstView[start + k] = srcView[k]; }
                );
            }
        }
        else if (std::holds_alternative<GPUExecutor>(exec_))
        {
            // GPU executor but not GPU-aware MPI: data was staged through host buffers.
            // Copy received host data back to the device in-place (via copy-assign, not
            // move-assign) so that value_.data() stays the same pointer and any View
            // obtained from value() before the exchange remains valid.
            auto valH = value_.copyToHost();
            for (const auto& buf : commBuffers_)
            {
                for (localIdx k = 0; k < buf.patchSize; k++)
                    valH.view()[buf.rangeStart + k] = buf.recvBuf[static_cast<std::size_t>(k)];
            }
            Vector<ValueType> backToDevice = valH.copyToExecutor(exec_);
            value_ = backToDevice; // operator=(const Vector&): setContainer dispatches async kernel
            // Fence before backToDevice goes out of scope: the setContainer kernel reads from
            // backToDevice.data_; freeing it while the kernel runs is a GPU page fault.
            fence(exec_);
        }
        else
        {
            // SerialExecutor or CPUExecutor: value_ lives in host-accessible memory;
            // write received data directly without any intermediate allocation.
            for (const auto& buf : commBuffers_)
            {
                ValueType* dst = value_.data() + buf.rangeStart;
                for (localIdx k = 0; k < buf.patchSize; k++)
                    dst[k] = buf.recvBuf[static_cast<std::size_t>(k)];
            }
        }
        requests_.clear();
        communicating_ = false;
        commBuffers_.clear();
#endif
    }
 
    std::pair<localIdx, localIdx> range(localIdx patchId) const
    {
        return {offset_.data()[patchId], offset_.data()[patchId + 1]};
    }
 
private:
 
    friend struct NoWaitAccess;
 
    Vector<ValueType>& valueNoWait() { return value_; }
 
    Executor exec_;                   
    mutable Vector<ValueType> value_; 
    Vector<ValueType> refValue_;      
    Vector<scalar>
        valueFraction_; 
    Vector<ValueType> refGrad_; 
    Vector<int> boundaryTypes_; 
    Vector<localIdx> offset_;   
    localIdx nBoundaries_;      
    localIdx nBoundaryFaces_;   
 
#ifdef NF_WITH_MPI_SUPPORT
    struct CommBuffer
    {
        std::vector<ValueType> sendBuf;                 // host staging, used when !gpuAwareMpi
        std::vector<ValueType> recvBuf;                 // host staging, used when !gpuAwareMpi
        std::optional<Vector<ValueType>> deviceRecvBuf; // device buffer, used when gpuAwareMpi
        localIdx rangeStart;
        localIdx patchSize;
    };
    mutable std::vector<MPI_Request>
        requests_; 
    mutable std::vector<CommBuffer>
        commBuffers_; 
    mutable bool communicating_ = false;
#endif
};
 
struct NoWaitAccess
{
    template<typename ValueType>
    static Vector<ValueType>& value(BoundaryData<ValueType>& in)
    {
        return in.valueNoWait();
    }
};
 
}