OnnxCompute.h

浏览该文件的文档.

#pragma once
#include <cstdint>
#include <memory>
#include <span>
#include <string>
#include <vector>
#include "common/Result.h"
#include "tensor/OnnxStorage.h"
 
namespace eve::tensor {
struct OnnxKernel {
    std::string                           source;
    std::vector<std::span<const uint8_t>> inputs;
    size_t                                outputBytes = 0;
    uint32_t                              workItems   = 0;
};
class OnnxCompute {
public:
    virtual ~OnnxCompute() = default;
    [[nodiscard]] virtual Result<std::vector<uint8_t>> dispatch(const OnnxKernel& kernel) = 0;
 
    [[nodiscard]] virtual Result<OnnxBuffer> enqueue(const std::string& source, std::span<const OnnxBuffer> inputs,
                                                     size_t outputBytes, uint32_t workItems) {
        OnnxKernel                        k{source, {}, outputBytes, workItems};
        std::vector<std::vector<uint8_t>> copies;
        copies.reserve(inputs.size());
        for (const auto& input : inputs) {
            if (input.host)
                k.inputs.push_back(*input.host);
            else {
                if (!input.device)
                    return Result<OnnxBuffer>::failure(
                        Diagnostic::error(DiagnosticCode::InvalidArgument, "Missing GPU input"));
                auto r = input.device->readback();
                if (!r.ok()) return Result<OnnxBuffer>::failure(r.status());
                copies.push_back(std::move(r.value()));
                k.inputs.push_back(copies.back());
            }
        }
        auto r = dispatch(k);
        if (!r.ok()) return Result<OnnxBuffer>::failure(r.status());
        auto bytes = std::make_shared<const std::vector<uint8_t>>(std::move(r.value()));
        return Result<OnnxBuffer>::success({bytes, {}, bytes->size()});
    }
    virtual OnnxTransferStats transferStats() const { return {}; }
 
protected:
    friend class OnnxModel;
    virtual void beginRun() noexcept {}
    virtual void endRun() noexcept {}
};
[[nodiscard]] EVENGINE_API_DOMAINS Result<std::unique_ptr<OnnxCompute>> createOnnxGpuCompute(
    uint32_t compilerWorkers = 4);
}  // namespace eve::tensor