map.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
// Copyright 2024 XCENA Inc.
 
#pragma once
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <type_traits>
#include <vector>
#include <xpti/xpti.hpp>
 
#include "pxl/ndarray.hpp"
#include "pxl/type.hpp"
 
namespace xpti
{
class Profile;
}
 
namespace pxl
{
class Function;
class Job;
class Stream;
 
namespace impl
{
class MapFactory;
class MapImpl;
}  // namespace impl
 
class Map
{
public:
    friend class impl::MapFactory;
 
    using CompletionCallback = std::function<void(void* arg)>;
 
    using MessageCallback = std::function<void(void* message, void* arg)>;
 
    using ErrorCallback = std::function<void(void* arg)>;
 
    void setCompletionCallback(const CompletionCallback& callback, void* arg);
 
    void setMessageCallback(const MessageCallback& callback, void* arg);
 
    void setErrorCallback(const ErrorCallback& callback, void* arg);
 
    void setBatchSize(const uint32_t& batchSize);
 
    uint32_t getBatchSize() const;
 
    uint32_t getTaskCount() const;
 
    void setClusterBitmap(const uint32_t& clusterBitmap);
 
    void setLocalityMode(const LocalityMode& mode);
 
    void setStream();
 
    void setStream(Stream* stream);
 
    uint32_t streamId() const;
 
    Result execute()
    {
        return Result::InvalidArgument;
    }
 
    template <typename... ARGS>
    Result execute(ARGS&&... args)
    {
        std::vector<ArgInfo_t> argInfos;
        if (buildArgInfo(argInfos, std::forward<ARGS>(args)...) == false)
        {
            return Result::InvalidArgument;
        }
 
        return execute(argInfos);
    }
 
    Result execute(std::vector<ArgInfo_t>& args);
 
    template <typename... ARG>
    Result setInput(ARG&&... args)
    {
        std::vector<ArgInfo_t> argInfos;
        if (buildArgInfo(argInfos, std::forward<ARG>(args)...) == false)
        {
            return Result::InvalidArgument;
        }
        return setInput(argInfos);
    }
 
    template <typename... ARG>
    Result setOutput(ARG&&... args)
    {
        std::vector<ArgInfo_t> argInfos;
        if (buildArgInfo(argInfos, std::forward<ARG>(args)...) == false)
        {
            return Result::InvalidArgument;
        }
        return setOutput(argInfos);
    }
 
    Result setInput(std::vector<ArgInfo_t>& args);
 
    Result setOutput(std::vector<ArgInfo_t>& args);
 
    Result cancel();
 
    Result synchronize();
 
    ExecuteStatus getExecuteStatus() const;
 
    Progress_t getProgress() const;
 
    std::vector<KernelError_t> getKernelError();
 
    std::string getStats() const;
 
    uint64_t id() const;
 
    uint32_t deviceId() const;
 
    Result enableProfiling(const xpti::ProfileConfig& config);
 
    template <xpti::QueryKey Key>
    xpti::SessionData<Key> getProfilingData(const xpti::QueryFilter& filter = {}) const
    {
        // Use this map's id() as sessionId
        return xpti::getSessionData<Key>(id(), filter);
    }
 
    Result exportProfiling(const std::string& path, xpti::ExportFormat format) const;
 
    Result exportProfilingCsv(const std::string& path = "") const;
 
protected:
    virtual ~Map() = default;
 
private:
    Map(std::unique_ptr<impl::MapImpl> impl);
 
    template <typename ARG, typename... ARGS>
    typename std::enable_if<std::is_trivial<typename std::remove_reference<ARG>::type>::value, bool>::type
    buildArgInfo(std::vector<ArgInfo_t>& argInfos, ARG&& arg, ARGS&&... args)
    {
        argInfos.emplace_back(ArgInfo_t{ArgType::Constant, &arg, sizeof(ARG)});
        if constexpr (sizeof...(args) > 0)
        {
            return buildArgInfo(argInfos, std::forward<ARGS>(args)...);
        }
        return true;
    }
 
    template <typename ARG, typename... ARGS>
    typename std::enable_if<!std::is_trivial<typename std::remove_reference<ARG>::type>::value, bool>::type
    buildArgInfo([[maybe_unused]] std::vector<ArgInfo_t>& argInfos, [[maybe_unused]] ARG&& arg, [[maybe_unused]] ARGS&&... args)
    {
        return false;
    }
 
    template <typename ARG, typename... ARGS>
    bool buildArgInfo(std::vector<ArgInfo_t>& argInfos, NDArray<ARG>& ndarrayArg, ARGS&&... args)
    {
        argInfos.emplace_back(ArgInfo_t{ArgType::NDArray, &ndarrayArg, sizeof(NDArray<ARG>)});
        if constexpr (sizeof...(args) > 0)
        {
            return buildArgInfo(argInfos, std::forward<ARGS>(args)...);
        }
        return true;
    }
 
    template <typename ARG, typename... ARGS>
    bool buildArgInfo(std::vector<ArgInfo_t>& argInfos, NDArray<ARG>&& ndarrayArg, ARGS&&... args) = delete;
 
    template <typename ARG, typename... ARGS>
    bool buildArgInfo(std::vector<ArgInfo_t>& argInfos, ARG* arg, ARGS&&... args)
    {
        argInfos.emplace_back(ArgInfo_t{ArgType::DeviceMemory, arg, sizeof(void*)});
        if constexpr (sizeof...(args) > 0)
        {
            return buildArgInfo(argInfos, std::forward<ARGS>(args)...);
        }
        return true;
    }
 
    std::unique_ptr<impl::MapImpl> impl_;
};
 
}  // namespace pxl