ort_session.cpp

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

#include "models/backend/ort_session.h"

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <string>
#include <unordered_map>

#include "glog/logging.h"

#include "common/file_path_util.h"
#include "models/backend/session_io.h"

namespace mortred {
namespace models {
namespace backend {

using mortred::common::StatusCode;

namespace {

ONNXTensorElementDataType to_ort_dtype(const DType& dtype) {
    switch (dtype) {
        case DType::F32:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
        case DType::I32:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32;
        case DType::I64:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64;
        case DType::U8:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8;
        default:
            return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
    }
}

bool dtype_from_ort(const ONNXTensorElementDataType& dtype, DType* out, std::string* err) {
    switch (dtype) {
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
            *out = DType::F32;
            return true;
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32:
            *out = DType::I32;
            return true;
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64:
            *out = DType::I64;
            return true;
        case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
            *out = DType::U8;
            return true;
        default:
            if (err != nullptr) {
                *err = "unsupported onnxruntime tensor element type: " + std::to_string(dtype);
            }
            return false;
    }
}

const void* ort_tensor_data(const Ort::Value& value, const DType& dtype) {
    switch (dtype) {
        case DType::F32:
            return value.GetTensorData<float>();
        case DType::I32:
            return value.GetTensorData<int32_t>();
        case DType::I64:
            return value.GetTensorData<int64_t>();
        case DType::U8:
            return value.GetTensorData<uint8_t>();
        default:
            return nullptr;
    }
}

}  // namespace

StatusCode OrtSession::init(const BackendConfig& config, std::string* err) {
    if (err != nullptr) {
        err->clear();
    }
    if (!mortred::common::FilePathUtil::is_file_exist(config.model_file_path)) {
        if (err != nullptr) {
            *err = "onnx model file not exist: " + config.model_file_path;
        }
        return StatusCode::MODEL_INIT_FAILED;
    }

    try {
        _m_env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, "mortred");
        _m_session_options = Ort::SessionOptions();
        _m_session_options.SetIntraOpNumThreads(config.threads);
        _m_session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
        _m_session_options.SetExecutionMode(ExecutionMode::ORT_SEQUENTIAL);
        if (config.use_gpu()) {
            int limit_mb = config.gpu_mem_limit_mb;
            if (const char* env = std::getenv("MORTRED_ORT_GPU_MEM_LIMIT_MB");
                env != nullptr && *env != '\0') {
                try {
                    limit_mb = std::stoi(env);
                } catch (...) {
                    if (err != nullptr) {
                        *err = "invalid MORTRED_ORT_GPU_MEM_LIMIT_MB";
                    }
                    return StatusCode::MODEL_INIT_FAILED;
                }
            }
            if (limit_mb < 0) {
                if (err != nullptr) {
                    *err = "onnxruntime gpu_mem_limit_mb must be >= 0 (0 = unlimited)";
                }
                return StatusCode::MODEL_INIT_FAILED;
            }
            size_t limit_bytes = 0;
            if (limit_mb > 0) {
                const uint64_t bytes =
                    static_cast<uint64_t>(limit_mb) * 1024ull * 1024ull;
                if (bytes > std::numeric_limits<size_t>::max()) {
                    if (err != nullptr) {
                        *err = "onnxruntime gpu_mem_limit_mb overflows size_t";
                    }
                    return StatusCode::MODEL_INIT_FAILED;
                }
                limit_bytes = static_cast<size_t>(bytes);
            }
            Ort::CUDAProviderOptions cuda_opts;
            const std::string mem_limit =
                (limit_mb == 0)
                    ? std::to_string(std::numeric_limits<size_t>::max())
                    : std::to_string(limit_bytes);
            cuda_opts.Update({
                {"device_id", std::to_string(config.device_id)},
                {"gpu_mem_limit", mem_limit},
                {"arena_extend_strategy", "kSameAsRequested"},
                {"cudnn_conv_algo_search", "DEFAULT"},
                {"do_copy_in_default_stream", "1"},
            });
            _m_session_options.AppendExecutionProvider_CUDA_V2(*cuda_opts);
            LOG(INFO) << "onnxruntime cuda gpu_mem_limit_mb=" << limit_mb
                      << (limit_mb == 0 ? " (unlimited)" : "")
                      << " device_id=" << config.device_id;
            // fuse/transform passes of ALL are not guaranteed on the CUDA EP
            _m_session_options.SetGraphOptimizationLevel(
                GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
        }
        _m_session = std::make_unique<Ort::Session>(
            _m_env, config.model_file_path.c_str(), _m_session_options);

        Ort::AllocatorWithDefaultOptions allocator;
        const auto fill_infos = [&allocator](size_t count, auto get_name, auto get_type_info,
                                             std::vector<TensorInfo>* infos,
                                             std::string* parse_err) -> bool {
            for (size_t idx = 0; idx < count; ++idx) {
                TensorInfo info;
                info.name = get_name(idx, allocator);
                // ConstTensorTypeAndShapeInfo borrows from the TypeInfo: keep
                // the owner alive while reading shape/element type
                const auto type_info = get_type_info(idx);
                const auto shape_and_type = type_info.GetTensorTypeAndShapeInfo();
                if (!dtype_from_ort(shape_and_type.GetElementType(), &info.dtype, parse_err)) {
                    return false;
                }
                info.shape = shape_and_type.GetShape();
                info.dynamic = shape_is_dynamic(info.shape);
                info.layout = host_output_layout(info.shape);
                infos->push_back(std::move(info));
            }
            return true;
        };

        const bool inputs_ok = fill_infos(
            _m_session->GetInputCount(),
            [this](size_t idx, Ort::AllocatorWithDefaultOptions& alloc) {
                return std::string(_m_session->GetInputNameAllocated(idx, alloc).get());
            },
            [this](size_t idx) { return _m_session->GetInputTypeInfo(idx); }, &_m_input_infos, err);
        if (!inputs_ok) {
            return StatusCode::MODEL_INIT_FAILED;
        }
        const bool outputs_ok = fill_infos(
            _m_session->GetOutputCount(),
            [this](size_t idx, Ort::AllocatorWithDefaultOptions& alloc) {
                return std::string(_m_session->GetOutputNameAllocated(idx, alloc).get());
            },
            [this](size_t idx) { return _m_session->GetOutputTypeInfo(idx); }, &_m_output_infos,
            err);
        if (!outputs_ok) {
            return StatusCode::MODEL_INIT_FAILED;
        }
    } catch (const Ort::Exception& exception) {
        if (err != nullptr) {
            *err = "onnxruntime init failed: " + std::string(exception.what());
        }
        _m_session.reset();
        return StatusCode::MODEL_INIT_FAILED;
    }

    if (_m_input_infos.empty() || _m_output_infos.empty()) {
        if (err != nullptr) {
            *err = "onnx model exposes no io tensors: " + config.model_file_path;
        }
        return StatusCode::MODEL_INIT_FAILED;
    }
    if (apply_configured_io_names(config.input_names, &_m_input_infos, "input", err) !=
        StatusCode::OK) {
        return StatusCode::MODEL_INIT_FAILED;
    }
    if (apply_configured_io_names(config.output_names, &_m_output_infos, "output", err) !=
        StatusCode::OK) {
        return StatusCode::MODEL_INIT_FAILED;
    }
    const auto rebuild_name_ptrs = [](const std::vector<TensorInfo>& infos,
                                      std::vector<std::string>* names,
                                      std::vector<const char*>* name_ptrs) {
        names->clear();
        names->reserve(infos.size());
        for (const auto& info : infos) {
            names->push_back(info.name);
        }
        name_ptrs->clear();
        name_ptrs->reserve(names->size());
        for (const auto& name : *names) {
            name_ptrs->push_back(name.c_str());
        }
    };
    rebuild_name_ptrs(_m_input_infos, &_m_input_names, &_m_input_name_ptrs);
    rebuild_name_ptrs(_m_output_infos, &_m_output_names, &_m_output_name_ptrs);

    _m_model_file_path = config.model_file_path;
    for (const auto& info : _m_input_infos) {
        LOG(INFO) << "ort session input: " << info.to_string();
    }
    for (const auto& info : _m_output_infos) {
        LOG(INFO) << "ort session output: " << info.to_string();
    }
    return StatusCode::OK;
}

StatusCode OrtSession::run(const std::vector<NamedTensor>& inputs,
                           std::vector<NamedTensor>& outputs) {
    if (_m_session == nullptr) {
        LOG(ERROR) << "onnxruntime session is not initialized";
        return StatusCode::MODEL_INIT_FAILED;
    }
    std::vector<const NamedTensor*> ordered_inputs;
    std::string match_err;
    if (match_required_inputs(_m_input_infos, inputs, &ordered_inputs, &match_err) !=
        StatusCode::OK) {
        LOG(ERROR) << "onnxruntime " << match_err;
        return StatusCode::MODEL_RUN_SESSION_FAILED;
    }

    try {
        const auto memory_info = Ort::MemoryInfo::CreateCpu(
            OrtAllocatorType::OrtDeviceAllocator, OrtMemType::OrtMemTypeDefault);
        std::vector<Ort::Value> ort_inputs;
        ort_inputs.reserve(_m_input_infos.size());
        // Bind in session input order so name ptrs and Ort::Value slots match,
        // independent of the caller's NamedTensor vector order.
        for (size_t idx = 0; idx < _m_input_infos.size(); ++idx) {
            const auto& info = _m_input_infos[idx];
            const auto& named = *ordered_inputs[idx];
            if (info.dtype != named.tensor.dtype) {
                LOG(ERROR) << "onnxruntime input '" << named.name << "' dtype mismatch, expected "
                           << info.to_string() << ", got "
                           << dtype_to_string(named.tensor.dtype);
                return StatusCode::MODEL_RUN_SESSION_FAILED;
            }
            if (!named.tensor.shape_is_concrete()) {
                LOG(ERROR) << "onnxruntime input '" << named.name << "' has non concrete shape "
                           << shape_to_string(named.tensor.shape);
                return StatusCode::MODEL_RUN_SESSION_FAILED;
            }
            ort_inputs.push_back(Ort::Value::CreateTensor(
                memory_info, const_cast<uint8_t*>(named.tensor.buffer.data()),
                named.tensor.byte_size(), named.tensor.shape.data(), named.tensor.shape.size(),
                to_ort_dtype(named.tensor.dtype)));
        }

        auto ort_outputs = _m_session->Run(
            Ort::RunOptions{nullptr}, _m_input_name_ptrs.data(), ort_inputs.data(),
            ort_inputs.size(), _m_output_name_ptrs.data(), _m_output_name_ptrs.size());

        std::string count_err;
        if (require_output_count(_m_output_infos.size(), ort_outputs.size(), "onnxruntime",
                                 &count_err) != StatusCode::OK) {
            LOG(ERROR) << count_err;
            return StatusCode::MODEL_RUN_SESSION_FAILED;
        }

        outputs.clear();
        outputs.reserve(ort_outputs.size());
        for (size_t idx = 0; idx < ort_outputs.size(); ++idx) {
            const auto& value = ort_outputs[idx];
            NamedTensor named;
            named.name = _m_output_infos[idx].name;
            const auto type_and_shape = value.GetTensorTypeAndShapeInfo();
            std::string dtype_err;
            if (!dtype_from_ort(type_and_shape.GetElementType(), &named.tensor.dtype,
                                &dtype_err)) {
                LOG(ERROR) << "onnxruntime output '" << named.name << "': " << dtype_err;
                return StatusCode::MODEL_RUN_SESSION_FAILED;
            }
            named.tensor.shape = type_and_shape.GetShape();
            size_t bytes = 0;
            if (!checked_shape_nbytes(named.tensor.shape, named.tensor.dtype, &bytes) ||
                bytes == 0) {
                LOG(ERROR) << "onnxruntime output '" << named.name
                           << "' is empty or shape byte size overflows: "
                           << shape_to_string(named.tensor.shape);
                return StatusCode::MODEL_EMPTY_OUTPUT;
            }
            named.tensor.buffer.resize(bytes);
            const void* src = ort_tensor_data(value, named.tensor.dtype);
            if (src == nullptr) {
                LOG(ERROR) << "onnxruntime output '" << named.name << "' data is null";
                return StatusCode::MODEL_RUN_SESSION_FAILED;
            }
            std::memcpy(named.tensor.buffer.data(), src, bytes);
            named.tensor.layout = host_output_layout(named.tensor.shape);
            outputs.push_back(std::move(named));
        }
    } catch (const Ort::Exception& exception) {
        LOG(ERROR) << "onnxruntime run failed: " << exception.what();
        return StatusCode::MODEL_RUN_SESSION_FAILED;
    }
    return StatusCode::OK;
}

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