session.h

/************************************************
* Copyright MaybeShewill-CV. All Rights Reserved.
* Author: MaybeShewill-CV
* File: session.h
* Date: 26-8-20
************************************************/

#ifndef MORTRED_MODELS_BACKEND_SESSION_H
#define MORTRED_MODELS_BACKEND_SESSION_H

#include <memory>
#include <string>
#include <vector>

#include "common/status_code.h"
#include "models/backend/backend_config.h"
#include "models/backend/tensor.h"

namespace mortred {
namespace models {
namespace backend {
using mortred::common::StatusCode;

/***
 * Metadata of one model io tensor. shape may contain -1 dims for dynamic
 * axes (dynamic == true); concrete shapes are required at run time.
 */
struct TensorInfo {
    std::string name;
    DType dtype = DType::F32;
    std::vector<int64_t> shape;
    bool dynamic = false;
    TensorLayout layout = TensorLayout::Unknown;

    std::string to_string() const {
        std::string out = name + ":";
        out += dtype_to_string(dtype);
        out += shape_to_string(shape);
        if (layout != TensorLayout::Unknown) {
            out += " ";
            out += tensor_layout_to_string(layout);
        }
        if (dynamic) {
            out += " (dynamic)";
        }
        return out;
    }
};

/***
 * Backend-agnostic inference session. One session per worker (sessions are
 * NOT required to be thread safe), all resources are owned through RAII.
 */
class InferenceSession {
  public:
    virtual ~InferenceSession() = default;

    InferenceSession() = default;
    InferenceSession(const InferenceSession&) = delete;
    InferenceSession& operator=(const InferenceSession&) = delete;

    /***
     * build a session from a parsed backend config. Returns nullptr and fills
     * err on failure; never throws for expected failures (bad file, unknown
     * io names, unsupported dtype).
     */
    static std::unique_ptr<InferenceSession> create(const BackendConfig& config,
                                                    std::string* err = nullptr);

    virtual const std::vector<TensorInfo>& inputs() const = 0;
    virtual const std::vector<TensorInfo>& outputs() const = 0;

    /***
     * run one inference. Input tensors are validated against the model io
     * (name / dtype / concrete shape); dynamic shapes are set per run. Output
     * buffers are reused inside the session, the returned tensors own a host
     * copy with the concrete run-time shapes. Rank-4 tensors are NCHW when C
     * is dim1 ([N,C,H,W]) and NHWC when C is dim3 ([N,H,W,C]); rank-1/2 are
     * Linear. Rank-3 stays Unknown.
     *
     * inputs()/outputs() follow [MODEL.backend] input_names and output_names
     * when those vectors are non-empty (subset + configured order). Empty
     * names keep backend-native discovery order.
     */
    virtual StatusCode run(const std::vector<NamedTensor>& inputs,
                                         std::vector<NamedTensor>& outputs) = 0;
};

}  // namespace backend
}  // namespace models
}  // namespace mortred

#endif  // MORTRED_MODELS_BACKEND_SESSION_H