dpu.hpp

/* Copyright 2020 UPMEM. All rights reserved.
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include <atomic>
#include <cstdarg>
#include <climits>
#include <functional>
#include <ostream>
#include <string>
#include <vector>

extern "C" {
#include <dpu.h>
#include <dpu_log_internals.h>
#include <dpu_management.h>
}

namespace dpu {

const unsigned ALLOCATE_ALL = DPU_ALLOCATE_ALL;

template <class F>
static bool
__get_block(struct sg_block_info *out, uint32_t dpu_index, uint32_t block_index, void *args)
{
    auto f = static_cast<F *>(args);
    return (*f)(out, dpu_index, block_index);
}

class DpuError : public std::exception {
    friend class DpuProgram;
    friend class DpuSetOps;
    friend class DpuSet;
    friend class DpuSetAsync;

public:
    virtual const char *
    what() const noexcept override
    {
        return msg;
    }

private:
    dpu_error_t errorId;
    const char *msg;

    explicit DpuError(dpu_error_t ErrorId)
        : errorId(ErrorId)
    {
        msg = dpu_error_to_string(errorId);
    }

    ~DpuError() { free((void *)msg); }

    static void
    throwOnErr(dpu_error_t Error)
    {
        if (Error != DPU_OK) {
            throw DpuError(Error);
        }
    }
};

class DpuSymbol {
    friend class DpuProgram;
    friend class DpuSetOps;

public:
    DpuSymbol(unsigned Address, unsigned Size)
        : cSymbol({ .address = Address, .size = Size })
    {
    }

private:
    DpuSymbol() { }
    struct dpu_symbol_t cSymbol;
};

class DpuProgram {
    friend class DpuSet;

public:
    DpuSymbol
    get(const std::string &SymbolName)
    {
        DpuSymbol symbol;
        DpuError::throwOnErr(dpu_get_symbol(cProgram, SymbolName.c_str(), &symbol.cSymbol));
        return symbol;
    }

private:
    struct dpu_program_t *cProgram { nullptr };
};

class DpuSet;
class DpuSetAsync;

using CallbackFn = std::function<void(DpuSet &, unsigned)>;

class DpuSetOps {
    friend class DpuSet;
    friend class DpuSetAsync;

public:
    template <typename T>
    void
    copy(const std::string &DstSymbol, unsigned Offset, const std::vector<T> &SrcBuffer, unsigned Size)
    {
        dpu_xfer_flags_t flags = async ? DPU_XFER_ASYNC : DPU_XFER_DEFAULT;
        DpuError::throwOnErr(dpu_broadcast_to(cSet, DstSymbol.c_str(), Offset, SrcBuffer.data(), Size, flags));
    }

    template <typename T>
    void
    copy(const std::string &DstSymbol, unsigned Offset, const std::vector<T> &SrcBuffer)
    {
        copy(DstSymbol, Offset, SrcBuffer, SrcBuffer.size() * sizeof(T));
    }

    template <typename T>
    void
    copy(const std::string &DstSymbol, const std::vector<T> &SrcBuffer, unsigned Size)
    {
        copy(DstSymbol, 0, SrcBuffer, Size);
    }

    template <typename T>
    void
    copy(const std::string &DstSymbol, const std::vector<T> &SrcBuffer)
    {
        copy(DstSymbol, 0, SrcBuffer, SrcBuffer.size() * sizeof(T));
    }

    template <typename T>
    void
    copy(DpuSymbol &DstSymbol, unsigned Offset, const std::vector<T> &SrcBuffer, unsigned Size)
    {
        dpu_xfer_flags_t flags = async ? DPU_XFER_ASYNC : DPU_XFER_DEFAULT;
        DpuError::throwOnErr(dpu_broadcast_to_symbol(cSet, DstSymbol.cSymbol, Offset, SrcBuffer.data(), Size, flags));
    }

    template <typename T>
    void
    copy(DpuSymbol &DstSymbol, unsigned Offset, const std::vector<T> &SrcBuffer)
    {
        copy(DstSymbol, Offset, SrcBuffer, SrcBuffer.size() * sizeof(T));
    }

    template <typename T>
    void
    copy(DpuSymbol &DstSymbol, const std::vector<T> &SrcBuffer, unsigned Size)
    {
        copy(DstSymbol, 0, SrcBuffer, Size);
    }

    template <typename T>
    void
    copy(DpuSymbol &DstSymbol, const std::vector<T> &SrcBuffer)
    {
        copy(DstSymbol, 0, SrcBuffer, SrcBuffer.size() * sizeof(T));
    }

    template <typename T>
    void
    copy(const std::string &DstSymbol, unsigned Offset, const std::vector<std::vector<T>> &SrcBuffers, unsigned Size)
    {
        struct dpu_set_t dpu;
        unsigned dpuIdx;

        DPU_FOREACH (cSet, dpu, dpuIdx) {
            DpuError::throwOnErr(dpu_prepare_xfer(dpu, (void *)SrcBuffers[dpuIdx].data()));
        }

        dpu_xfer_flags_t flags = async ? DPU_XFER_ASYNC : DPU_XFER_DEFAULT;
        DpuError::throwOnErr(dpu_push_xfer(cSet, DPU_XFER_TO_DPU, DstSymbol.c_str(), Offset, Size, flags));
    }

    template <typename T>
    void
    copy(const std::string &DstSymbol, unsigned Offset, const std::vector<std::vector<T>> &SrcBuffers)
    {
        if (SrcBuffers.size() == 0) {
            DpuError::throwOnErr(DPU_ERR_INVALID_MEMORY_TRANSFER);
        }

        unsigned nrElements = SrcBuffers[0].size();
        for (auto buf : SrcBuffers) {
            if (nrElements != buf.size()) {
                DpuError::throwOnErr(DPU_ERR_INVALID_MEMORY_TRANSFER);
            }
        }

        copy(DstSymbol, Offset, SrcBuffers, nrElements * sizeof(T));
    }

    template <typename T>
    void
    copy(const std::string &DstSymbol, const std::vector<std::vector<T>> &SrcBuffers, unsigned Size)
    {
        copy(DstSymbol, 0, SrcBuffers, Size);
    }

    template <typename T>
    void
    copy(const std::string &DstSymbol, const std::vector<std::vector<T>> &SrcBuffers)
    {
        copy(DstSymbol, 0, SrcBuffers);
    }

    void
    copyScatterGather(const std::string &DstSymbol,
        unsigned Offset,
        get_block_t &get_block_info,
        unsigned Size,
        bool length_check = true)
    {
        dpu_sg_xfer_flags_t flags = async ? DPU_SG_XFER_ASYNC : DPU_SG_XFER_DEFAULT;
        if (!length_check) {
            flags = static_cast<dpu_sg_xfer_flags_t>(flags | DPU_SG_XFER_DISABLE_LENGTH_CHECK);
        }
        DpuError::throwOnErr(dpu_push_sg_xfer(cSet, DPU_XFER_TO_DPU, DstSymbol.c_str(), Offset, Size, &get_block_info, flags));
    }

    void
    copyScatterGather(const std::string &DstSymbol, get_block_t &get_block_info, unsigned Size, bool length_check = true)
    {
        copyScatterGather(DstSymbol, 0, get_block_info, Size, length_check);
    }

    template <class F>
    void
    copyScatterGather(const std::string &DstSymbol, unsigned Offset, F f, unsigned Size, bool length_check = true)
    {
        get_block_t get_block_info { __get_block<F>, &f, sizeof(f) };
        copyScatterGather(DstSymbol, Offset, get_block_info, Size, length_check);
    }

    template <class F>
    void
    copyScatterGather(const std::string &DstSymbol, F f, unsigned Size, bool length_check = true)
    {
        copyScatterGather(DstSymbol, 0, f, Size, length_check);
    }

    template <typename T>
    void
    copy(DpuSymbol &DstSymbol, unsigned Offset, const std::vector<std::vector<T>> &SrcBuffers, unsigned Size)
    {
        struct dpu_set_t dpu;
        unsigned dpuIdx;

        DPU_FOREACH (cSet, dpu, dpuIdx) {
            DpuError::throwOnErr(dpu_prepare_xfer(dpu, (void *)SrcBuffers[dpuIdx].data()));
        }

        dpu_xfer_flags_t flags = async ? DPU_XFER_ASYNC : DPU_XFER_DEFAULT;
        DpuError::throwOnErr(dpu_push_xfer_symbol(cSet, DPU_XFER_TO_DPU, DstSymbol.cSymbol, Offset, Size, flags));
    }

    template <typename T>
    void
    copy(DpuSymbol &DstSymbol, unsigned Offset, const std::vector<std::vector<T>> &SrcBuffers)
    {
        if (SrcBuffers.size() == 0) {
            DpuError::throwOnErr(DPU_ERR_INVALID_MEMORY_TRANSFER);
        }

        unsigned nrElements = SrcBuffers[0].size();
        for (auto buf : SrcBuffers) {
            if (nrElements != buf.size()) {
                DpuError::throwOnErr(DPU_ERR_INVALID_MEMORY_TRANSFER);
            }
        }

        copy(DstSymbol, Offset, SrcBuffers, nrElements * sizeof(T));
    }

    template <typename T>
    void
    copy(DpuSymbol &DstSymbol, const std::vector<std::vector<T>> &SrcBuffers, unsigned Size)
    {
        copy(DstSymbol, 0, SrcBuffers, Size);
    }

    template <typename T>
    void
    copy(DpuSymbol &DstSymbol, const std::vector<std::vector<T>> &SrcBuffers)
    {
        copy(DstSymbol, 0, SrcBuffers);
    }

    void
    copyScatterGather(DpuSymbol &DstSymbol, unsigned Offset, get_block_t &get_block_info, unsigned Size, bool length_check = true)
    {
        dpu_sg_xfer_flags_t flags = async ? DPU_SG_XFER_ASYNC : DPU_SG_XFER_DEFAULT;
        if (!length_check) {
            flags = static_cast<dpu_sg_xfer_flags_t>(flags | DPU_SG_XFER_DISABLE_LENGTH_CHECK);
        }
        DpuError::throwOnErr(
            dpu_push_sg_xfer_symbol(cSet, DPU_XFER_TO_DPU, DstSymbol.cSymbol, Offset, Size, &get_block_info, flags));
    }

    void
    copyScatterGather(DpuSymbol &DstSymbol, get_block_t &get_block_info, unsigned Size, bool length_check = true)
    {
        copyScatterGather(DstSymbol, 0, get_block_info, Size, length_check);
    }

    template <class F>
    void
    copyScatterGather(DpuSymbol &DstSymbol, unsigned Offset, F f, unsigned Size, bool length_check = true)
    {
        get_block_t get_block_info { __get_block<F>, &f, sizeof(f) };
        copyScatterGather(DstSymbol, Offset, get_block_info, Size, length_check);
    }

    template <class F>
    void
    copyScatterGather(DpuSymbol &DstSymbol, F f, unsigned Size, bool length_check = true)
    {
        copyScatterGather(DstSymbol, 0, f, Size, length_check);
    }

    template <typename T>
    void
    copy(std::vector<std::vector<T>> &DstBuffers, unsigned Size, const std::string &SrcSymbol, unsigned Offset)
    {
        struct dpu_set_t dpu;
        unsigned dpuIdx;

        DPU_FOREACH (cSet, dpu, dpuIdx) {
            DpuError::throwOnErr(dpu_prepare_xfer(dpu, DstBuffers[dpuIdx].data()));
        }

        dpu_xfer_flags_t flags = async ? DPU_XFER_ASYNC : DPU_XFER_DEFAULT;
        DpuError::throwOnErr(dpu_push_xfer(cSet, DPU_XFER_FROM_DPU, SrcSymbol.c_str(), Offset, Size, flags));
    }

    template <typename T>
    void
    copy(std::vector<std::vector<T>> &DstBuffers, unsigned Size, const std::string &SrcSymbol)
    {
        copy(DstBuffers, Size, SrcSymbol, 0);
    }

    template <typename T>
    void
    copy(std::vector<std::vector<T>> &DstBuffers, const std::string &SrcSymbol, unsigned Offset)
    {
        if (DstBuffers.size() == 0) {
            DpuError::throwOnErr(DPU_ERR_INVALID_MEMORY_TRANSFER);
        }

        unsigned nrElements = DstBuffers[0].size();
        for (auto buf : DstBuffers) {
            if (nrElements != buf.size()) {
                DpuError::throwOnErr(DPU_ERR_INVALID_MEMORY_TRANSFER);
            }
        }

        copy(DstBuffers, nrElements * sizeof(T), SrcSymbol, Offset);
    }

    template <typename T>
    void
    copy(std::vector<std::vector<T>> &DstBuffers, const std::string &SrcSymbol)
    {
        copy(DstBuffers, SrcSymbol, 0);
    }

    void
    copyScatterGather(get_block_t &get_block_info,
        unsigned Size,
        const std::string &SrcSymbol,
        unsigned Offset,
        bool length_check = true)
    {
        dpu_sg_xfer_flags_t flags = async ? DPU_SG_XFER_ASYNC : DPU_SG_XFER_DEFAULT;
        if (!length_check) {
            flags = static_cast<dpu_sg_xfer_flags_t>(flags | DPU_SG_XFER_DISABLE_LENGTH_CHECK);
        }
        DpuError::throwOnErr(dpu_push_sg_xfer(cSet, DPU_XFER_FROM_DPU, SrcSymbol.c_str(), Offset, Size, &get_block_info, flags));
    }

    void
    copyScatterGather(get_block_t &get_block_info, unsigned Size, const std::string &SrcSymbol, bool length_check = true)
    {
        copyScatterGather(get_block_info, Size, SrcSymbol, 0, length_check);
    }

    template <class F>
    void
    copyScatterGather(F f, unsigned Size, const std::string &SrcSymbol, unsigned Offset, bool length_check = true)
    {
        get_block_t get_block_info { __get_block<F>, &f, sizeof(f) };
        copyScatterGather(get_block_info, Size, SrcSymbol, Offset, length_check);
    }

    template <class F>
    void
    copyScatterGather(F f, unsigned Size, const std::string &SrcSymbol, bool length_check = true)
    {
        copyScatterGather(f, Size, SrcSymbol, 0, length_check);
    }

    template <typename T>
    void
    copy(std::vector<std::vector<T>> &DstBuffers, unsigned Size, DpuSymbol &SrcSymbol, unsigned Offset)
    {
        struct dpu_set_t dpu;
        unsigned dpuIdx;

        DPU_FOREACH (cSet, dpu, dpuIdx) {
            DpuError::throwOnErr(dpu_prepare_xfer(dpu, DstBuffers[dpuIdx].data()));
        }

        dpu_xfer_flags_t flags = async ? DPU_XFER_ASYNC : DPU_XFER_DEFAULT;
        DpuError::throwOnErr(dpu_push_xfer_symbol(cSet, DPU_XFER_FROM_DPU, SrcSymbol.cSymbol, Offset, Size, flags));
    }

    template <typename T>
    void
    copy(std::vector<std::vector<T>> &DstBuffers, unsigned Size, DpuSymbol &SrcSymbol)
    {
        copy(DstBuffers, Size, SrcSymbol, 0);
    }

    template <typename T>
    void
    copy(std::vector<std::vector<T>> &DstBuffers, DpuSymbol &SrcSymbol, unsigned Offset)
    {
        if (DstBuffers.size() == 0) {
            DpuError::throwOnErr(DPU_ERR_INVALID_MEMORY_TRANSFER);
        }

        unsigned nrElements = DstBuffers[0].size();
        for (auto buf : DstBuffers) {
            if (nrElements != buf.size()) {
                DpuError::throwOnErr(DPU_ERR_INVALID_MEMORY_TRANSFER);
            }
        }

        copy(DstBuffers, nrElements * sizeof(T), SrcSymbol, Offset);
    }

    template <typename T>
    void
    copy(std::vector<std::vector<T>> &DstBuffers, DpuSymbol &SrcSymbol)
    {
        copy(DstBuffers, SrcSymbol, 0);
    }

    void
    copyScatterGather(get_block_t &get_block_info, unsigned Size, DpuSymbol &SrcSymbol, unsigned Offset, bool length_check = true)
    {
        dpu_sg_xfer_flags_t flags = async ? DPU_SG_XFER_ASYNC : DPU_SG_XFER_DEFAULT;
        if (!length_check) {
            flags = static_cast<dpu_sg_xfer_flags_t>(flags | DPU_SG_XFER_DISABLE_LENGTH_CHECK);
        }
        DpuError::throwOnErr(
            dpu_push_sg_xfer_symbol(cSet, DPU_XFER_FROM_DPU, SrcSymbol.cSymbol, Offset, Size, &get_block_info, flags));
    }

    void
    copyScatterGather(get_block_t &get_block_info, unsigned Size, DpuSymbol &SrcSymbol, bool length_check = true)
    {
        copyScatterGather(get_block_info, Size, SrcSymbol, 0, length_check);
    }

    template <class F>
    void
    copyScatterGather(F f, unsigned Size, DpuSymbol &SrcSymbol, unsigned Offset, bool length_check = true)
    {
        get_block_t get_block_info { __get_block<F>, &f, sizeof(f) };
        copyScatterGather(get_block_info, Size, SrcSymbol, Offset, length_check);
    }

    template <class F>
    void
    copyScatterGather(F f, unsigned Size, DpuSymbol &SrcSymbol, bool length_check = true)
    {
        copyScatterGather(f, Size, SrcSymbol, 0, length_check);
    }

    void
    exec()
    {
        dpu_launch_policy_t policy = async ? DPU_ASYNCHRONOUS : DPU_SYNCHRONOUS;
        DpuError::throwOnErr(dpu_launch(cSet, policy));
    }

private:
    struct dpu_set_t cSet;
    bool async;

    DpuSetOps(const struct dpu_set_t &CSet, bool Async)
        : cSet(CSet)
        , async(Async)
    {
    }
};

class DpuSet : public DpuSetOps {
    friend class DpuSetOps;
    friend class DpuSetAsync;

public:
    ~DpuSet()
    {
        if (manageCSet) {
            for (auto rank : _ranks) {
                delete rank;
            }
            for (auto dpu : _dpus) {
                delete dpu;
            }
            dpu_free(cSet);
        }
    }

    std::vector<DpuSet *> &
    dpus()
    {
        return _dpus;
    }

    std::vector<DpuSet *> &
    ranks()
    {
        return _ranks;
    }

    static DpuSet
    allocate(unsigned NrDpus = ALLOCATE_ALL, const std::string &Profile = "")
    {
        struct dpu_set_t cSet;
        DpuError::throwOnErr(dpu_alloc(NrDpus, Profile.c_str(), &cSet));
        return DpuSet(cSet);
    }

    static DpuSet
    allocateRanks(unsigned NrRanks = ALLOCATE_ALL, const std::string &Profile = "")
    {
        struct dpu_set_t cSet;
        DpuError::throwOnErr(dpu_alloc_ranks(NrRanks, Profile.c_str(), &cSet));
        return DpuSet(cSet);
    }

    DpuProgram
    load(const std::string &Executable)
    {
        DpuProgram Program;
        DpuError::throwOnErr(dpu_load(cSet, Executable.c_str(), &Program.cProgram));
        return Program;
    }

    void
    log(std::ostream &LogStream)
    {
        for (DpuSet *dpuSet : _dpus) {
            struct dpu_t *dpu = dpu_from_set(dpuSet->cSet);
            DpuError::throwOnErr(ostreamPrint(&LogStream, DPU_LOG_FORMAT_HEADER));
            DpuError::throwOnErr(dpulog_read_for_dpu_(dpu, ostreamPrint, &LogStream));
        }
    }

    DpuSetAsync
    async();

private:
    bool manageCSet;
    std::vector<DpuSet *> _dpus;
    std::vector<DpuSet *> _ranks;

    DpuSet(struct dpu_set_t CSet, bool ManageCSet = true, bool DetectChildren = true)
        : DpuSetOps(CSet, false)
        , manageCSet(ManageCSet)
    {

        if (DetectChildren) {
            struct dpu_set_t cRank;
            DPU_RANK_FOREACH (CSet, cRank) {
                DpuSet *rank = new DpuSet(cRank, false, false);
                struct dpu_set_t cDpu;

                DPU_FOREACH (cRank, cDpu) {
                    DpuSet *dpu = new DpuSet(cDpu, false, false);

                    dpu->_dpus.push_back(dpu);
                    dpu->_ranks.push_back(rank);
                    rank->_dpus.push_back(dpu);
                    _dpus.push_back(dpu);
                }

                rank->_ranks.push_back(rank);
                _ranks.push_back(rank);
            }
        }
    }

    static dpu_error_t
    ostreamPrint(void *Arg, const char *Fmt, ...)
    {
        std::ostream *LogStream = (std::ostream *)Arg;
        char *str;
        va_list ap;
        va_start(ap, Fmt);
        if (vasprintf(&str, Fmt, ap) == -1) {
            va_end(ap);
            return DPU_ERR_SYSTEM;
        }

        *LogStream << str;

        free(str);
        return DPU_OK;
    }
};

class DpuSetAsync : public DpuSetOps {
    friend class DpuSet;

    struct CallContext {
        CallbackFn callback;
        std::atomic_uint count;
    };

public:
    void
    call(const CallbackFn &Callback, bool IsBlocking, bool SingleCall)
    {
        unsigned flags = DPU_CALLBACK_ASYNC;
        if (!IsBlocking) {
            flags |= DPU_CALLBACK_NONBLOCKING;
        }
        if (SingleCall) {
            flags |= DPU_CALLBACK_SINGLE_CALL;
        }

        CallContext *context = new CallContext;
        context->callback = Callback;
        context->count = SingleCall ? 1 : set->_ranks.size();

        DpuError::throwOnErr(dpu_callback(cSet, cbWrapper, (void *)context, (dpu_callback_flags_t)flags));
    }

    void
    call(const CallbackFn &Callback)
    {
        call(Callback, true, false);
    }

    void
    sync()
    {
        DpuError::throwOnErr(dpu_sync(set->cSet));
    }

private:
    DpuSet *set;

    explicit DpuSetAsync(DpuSet *Set)
        : DpuSetOps(Set->cSet, true)
        , set(Set)
    {
    }

    static dpu_error_t
    cbWrapper(struct dpu_set_t CSet, unsigned Idx, void *Arg)
    {
        DpuSet dpuSet(CSet, false, true);
        CallContext *context = static_cast<CallContext *>(Arg);
        context->callback(dpuSet, Idx);
        if (--context->count == 0) {
            delete context;
        }
        return DPU_OK;
    }
};

inline DpuSetAsync
DpuSet::async()
{
    return DpuSetAsync(this);
}

}