onnxruntime/winml/lib/Api.Ort/OnnxruntimeModel.cpp

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "lib/Api.Ort/pch.h"
#include "OnnxruntimeModel.h"
#include "core/platform/windows/TraceLoggingConfig.h"
#include <evntrace.h>

#include "OnnxruntimeDescriptorConverter.h"
#include "OnnxruntimeEngine.h"
#include "OnnxruntimeErrors.h"

using namespace _winml;

struct winml_adapter_api_model_feature_helper {
  decltype(WinmlAdapterApi::ModelGetInputCount) GetCount;
  decltype(WinmlAdapterApi::ModelGetInputName) GetName;
  decltype(WinmlAdapterApi::ModelGetInputDescription) GetDescription;
  decltype(WinmlAdapterApi::ModelGetInputTypeInfo) GetTypeInfo;
};

HRESULT CreateFeatureDescriptors(
    OnnxruntimeEngineFactory* engine_factory,
    const winml_adapter_api_model_feature_helper* feature_helpers,
    OrtModel* ort_model,
    std::vector<OnnxruntimeValueInfoWrapper>& descriptors) {
  const auto ort_api = engine_factory->UseOrtApi();
  size_t count;
  RETURN_HR_IF_NOT_OK_MSG(feature_helpers->GetCount(ort_model, &count),
                          engine_factory->UseOrtApi());

  for (size_t i = 0; i < count; i++) {
    OnnxruntimeValueInfoWrapper descriptor;
    RETURN_HR_IF_NOT_OK_MSG(feature_helpers->GetName(ort_model, i, &descriptor.name_, &descriptor.name_length_),
                            engine_factory->UseOrtApi());
    RETURN_HR_IF_NOT_OK_MSG(feature_helpers->GetDescription(ort_model, i, &descriptor.description_, &descriptor.description_length_),
                            engine_factory->UseOrtApi());

    OrtTypeInfo* type_info;
    RETURN_HR_IF_NOT_OK_MSG(feature_helpers->GetTypeInfo(ort_model, i, &type_info),
                            engine_factory->UseOrtApi());

    descriptor.type_info_ = UniqueOrtTypeInfo(type_info, ort_api->ReleaseTypeInfo);

    descriptors.push_back(std::move(descriptor));
  }
  return S_OK;
}

HRESULT ModelInfo::RuntimeClassInitialize(_In_ OnnxruntimeEngineFactory* engine_factory, _In_ OrtModel* ort_model) {
  RETURN_HR_IF_NULL(E_INVALIDARG, ort_model);

  const auto winml_adapter_api = engine_factory->UseWinmlAdapterApi();

  // Get Metadata
  size_t count;
  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelGetMetadataCount(ort_model, &count),
                          engine_factory->UseOrtApi());

  const char* metadata_key;
  size_t metadata_key_len;
  const char* metadata_value;
  size_t metadata_value_len;
  for (size_t i = 0; i < count; i++) {
    RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelGetMetadata(ort_model, i, &metadata_key, &metadata_key_len, &metadata_value, &metadata_value_len),
                            engine_factory->UseOrtApi());

    model_metadata_.insert_or_assign(
        std::string(metadata_key, metadata_key_len),
        std::string(metadata_value, metadata_value_len));
  }

  _winml::OnnxruntimeDescriptorConverter converter(engine_factory, model_metadata_);

  static const winml_adapter_api_model_feature_helper input_helpers = {
      winml_adapter_api->ModelGetInputCount,
      winml_adapter_api->ModelGetInputName,
      winml_adapter_api->ModelGetInputDescription,
      winml_adapter_api->ModelGetInputTypeInfo};

  // Create inputs
  std::vector<OnnxruntimeValueInfoWrapper> inputs;
  RETURN_IF_FAILED(CreateFeatureDescriptors(engine_factory, &input_helpers, ort_model, inputs));
  input_features_ = converter.ConvertToLearningModelDescriptors(inputs.data(), inputs.size());

  // Create outputs
  static const winml_adapter_api_model_feature_helper output_helpers = {
      winml_adapter_api->ModelGetOutputCount,
      winml_adapter_api->ModelGetOutputName,
      winml_adapter_api->ModelGetOutputDescription,
      winml_adapter_api->ModelGetOutputTypeInfo};

  std::vector<OnnxruntimeValueInfoWrapper> outputs;
  RETURN_IF_FAILED(CreateFeatureDescriptors(engine_factory, &output_helpers, ort_model, outputs));
  output_features_ = converter.ConvertToLearningModelDescriptors(outputs.data(), outputs.size());

  const char* out;
  size_t len;

  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelGetAuthor(ort_model, &out, &len),
                          engine_factory->UseOrtApi());
  author_ = std::string(out, len);

  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelGetName(ort_model, &out, &len),
                          engine_factory->UseOrtApi());
  name_ = std::string(out, len);

  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelGetDomain(ort_model, &out, &len),
                          engine_factory->UseOrtApi());
  domain_ = std::string(out, len);

  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelGetDescription(ort_model, &out, &len),
                          engine_factory->UseOrtApi());
  description_ = std::string(out, len);

  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelGetVersion(ort_model, &version_),
                          engine_factory->UseOrtApi());
  return S_OK;
}

STDMETHODIMP ModelInfo::GetAuthor(const char** out, size_t* len) {
  *out = author_.c_str();
  *len = author_.size();
  return S_OK;
}

STDMETHODIMP ModelInfo::GetName(const char** out, size_t* len) {
  *out = name_.c_str();
  *len = name_.size();
  return S_OK;
}

STDMETHODIMP ModelInfo::SetName(const char* name) {
  name_ = std::string(name);
  return S_OK;
}

STDMETHODIMP ModelInfo::GetDomain(const char** out, size_t* len) {
  *out = domain_.c_str();
  *len = domain_.size();
  return S_OK;
}

STDMETHODIMP ModelInfo::GetDescription(const char** out, size_t* len) {
  *out = description_.c_str();
  *len = description_.size();
  return S_OK;
}

STDMETHODIMP ModelInfo::GetVersion(int64_t* out) {
  *out = version_;
  return S_OK;
}

struct CaseInsensitiveHash {
  size_t operator()(const winrt::hstring& key) const {
    size_t h = 0, i = 0;
    std::for_each(key.begin(), key.end(), [&](wchar_t c) {
      i++;
      h += i * towlower(c);
    });
    return h;
  }
};

struct CaseInsensitiveEqual {
  bool operator()(const winrt::hstring& left, const winrt::hstring& right) const {
    return left.size() == right.size() && std::equal(left.begin(), left.end(), right.begin(),
                                                     [](wchar_t a, wchar_t b) {
                                                       return towlower(a) == towlower(b);
                                                     });
  }
};

STDMETHODIMP ModelInfo::GetModelMetadata(ABI::Windows::Foundation::Collections::IMapView<HSTRING, HSTRING>** metadata) {
  std::unordered_map<winrt::hstring, winrt::hstring, CaseInsensitiveHash, CaseInsensitiveEqual> map_copy;
  for (auto& pair : model_metadata_) {
    auto metadata_key = _winml::Strings::HStringFromUTF8(pair.first);
    auto metadata_value = _winml::Strings::HStringFromUTF8(pair.second);
    map_copy.emplace(std::move(metadata_key), std::move(metadata_value));
  }
  auto map = winrt::single_threaded_map<winrt::hstring, winrt::hstring>(std::move(map_copy));
  winrt::copy_to_abi(map, *(void**)metadata);
  return S_OK;
}

STDMETHODIMP ModelInfo::GetInputFeatures(ABI::Windows::Foundation::Collections::IVectorView<winml::ILearningModelFeatureDescriptor>** features) {
  *features = nullptr;
  winrt::copy_to_abi(input_features_.GetView(), *(void**)features);
  return S_OK;
}

STDMETHODIMP ModelInfo::GetOutputFeatures(ABI::Windows::Foundation::Collections::IVectorView<winml::ILearningModelFeatureDescriptor>** features) {
  *features = nullptr;
  winrt::copy_to_abi(output_features_.GetView(), *(void**)features);
  return S_OK;
}

OnnruntimeModel::OnnruntimeModel() : ort_model_(nullptr, nullptr) {
}

HRESULT OnnruntimeModel::RuntimeClassInitialize(OnnxruntimeEngineFactory* engine_factory, UniqueOrtModel&& ort_model) {
  RETURN_HR_IF_NULL(E_INVALIDARG, ort_model);

  engine_factory_ = engine_factory;
  ort_model_ = std::move(ort_model);

  return S_OK;
}

STDMETHODIMP OnnruntimeModel::GetModelInfo(IModelInfo** info) {
  if (info_ == nullptr) {
    RETURN_IF_FAILED(Microsoft::WRL::MakeAndInitialize<ModelInfo>(&info_, engine_factory_.Get(), ort_model_.get()));
  }

  info_.CopyTo(info);

  return S_OK;
}

STDMETHODIMP OnnruntimeModel::ModelEnsureNoFloat16() {
  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();
  if (auto status = winml_adapter_api->ModelEnsureNoFloat16(ort_model_.get())) {
    return DXGI_ERROR_UNSUPPORTED;
  }
  return S_OK;
}

STDMETHODIMP OnnruntimeModel::CloneModel(IModel** copy) {
  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();

  OrtModel* ort_model_copy;
  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->CloneModel(ort_model_.get(), &ort_model_copy),
                          engine_factory_->UseOrtApi());

  auto model = UniqueOrtModel(ort_model_copy, winml_adapter_api->ReleaseModel);
  RETURN_IF_FAILED(Microsoft::WRL::MakeAndInitialize<OnnruntimeModel>(copy, engine_factory_.Get(), std::move(model)));

  return S_OK;
}

STDMETHODIMP OnnruntimeModel::SaveModel(_In_ const wchar_t* const file_name, _In_ unsigned size) {
  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();

  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->SaveModel(ort_model_.get(), file_name, size),
                          engine_factory_->UseOrtApi());
  return S_OK;
}

STDMETHODIMP OnnruntimeModel::SetName(const char* name) {
  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();
  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelSetName(ort_model_.get(), name),
                          engine_factory_->UseOrtApi());
  info_->SetName(name);
  return S_OK;
}

STDMETHODIMP OnnruntimeModel::DetachOrtModel(OrtModel** model) {
  *model = ort_model_.release();
  return S_OK;
}

HRESULT GetValue(const char* key, const char* const* keys, const char* const* values,
                 size_t num_values_in_dictionary, const char** value) {
  auto found_it =
      std::find_if(keys, keys + num_values_in_dictionary, [key](auto& key_name) {
        return _stricmp(key, key_name) == 0;
      });
  if (found_it == (keys + num_values_in_dictionary)) {
    return S_FALSE;
  }
  *value = values[std::distance(keys, found_it)];
  return S_OK;
}

STDMETHODIMP OnnruntimeModel::AddOperator(
    _In_ const char* const op_type, _In_ const char* const op_name, _In_ const char* const op_domain,
    _In_ const char* const* op_input_names, _In_ const char* const* actual_input_names, size_t num_inputs,
    _In_ const char* const* op_output_names, _In_ const char* const* actual_output_names, size_t num_outputs,
    _In_ const char* const* op_attribute_names, _In_ IValue** attribute_values, size_t num_attributes) {
  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();
  auto ort_api = engine_factory_->UseOrtApi();

  int32_t onnx_opset_version;
  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelGetOpsetVersion(ort_model_.get(), op_domain, &onnx_opset_version),
                          ort_api);
  size_t input_count;
  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->OperatorGetNumInputs(op_type, onnx_opset_version, op_domain, &input_count),
                           ort_api);

  size_t output_count;
  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->OperatorGetNumOutputs(op_type, onnx_opset_version, op_domain, &output_count),
                           ort_api);

  std::vector<const char*> input_names(input_count);
  for (size_t i = 0; i < input_count; i++) {
    const char* name;
    RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->OperatorGetInputName(op_type, onnx_opset_version, op_domain, i, &name),
                            ort_api);

    const char* actual_name;
    if (S_OK == GetValue(name, op_input_names, actual_input_names, num_inputs, &actual_name))
    {
      input_names[i] = actual_name;
    }
  }

  std::vector<const char*> output_names(output_count);
  for (size_t i = 0; i < output_count; i++) {
    const char* name;
    RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->OperatorGetOutputName(op_type, onnx_opset_version, op_domain, i, &name),
                             ort_api);
    const char* actual_name = nullptr;
    if (S_OK == GetValue(name, op_output_names, actual_output_names, num_outputs, &actual_name)) {
        output_names[i] = actual_name;
    }
  }

  std::vector<OrtValue*> attributes;
  for (size_t i = 0; i < num_attributes; i++) {
    attributes.push_back(static_cast<OnnxruntimeValue*>(*(attribute_values + i))->UseOrtValue());
  }

  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelAddOperator(
    ort_model_.get(), op_type, op_name, onnx_opset_version, op_domain, input_names.data(), input_count, output_names.data(), output_count, op_attribute_names, attributes.data(), num_attributes),
                          engine_factory_->UseOrtApi());
  return S_OK;
}

STDMETHODIMP OnnruntimeModel::AddModelInput(_In_ const char* const name, _In_ IDescriptorInfoProvider* descriptor_provider, bool is_constant, IValue* constant_value) {
  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();
  auto ort_api = engine_factory_->UseOrtApi();

  winrt::com_ptr<_winml::IDescriptorInfo> descriptor_info;
  descriptor_provider->GetDescriptorInfo(engine_factory_.Get(), descriptor_info.put());

  auto ort_type_info_provider = descriptor_info.as<_winml::IOrtTypeInfoProvider>();
  OrtTypeInfo* type_info;
  ort_type_info_provider->GetTypeInfo(&type_info);

  if (is_constant) {
    RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelAddConstantInput(ort_model_.get(), name, type_info, static_cast<OnnxruntimeValue*>(constant_value)->UseOrtValue()),
                            ort_api);
  } else {
    RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelAddInput(ort_model_.get(), name, type_info),
                            ort_api);
  }

  return S_OK;
}

STDMETHODIMP OnnruntimeModel::AddModelOutput(_In_ const char* const name, _In_ IDescriptorInfoProvider* descriptor_provider) {
  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();
  auto ort_api = engine_factory_->UseOrtApi();

  winrt::com_ptr<_winml::IDescriptorInfo> descriptor_info;
  descriptor_provider->GetDescriptorInfo(engine_factory_.Get(), descriptor_info.put());

  auto ort_type_info_provider = descriptor_info.as<_winml::IOrtTypeInfoProvider>();
  OrtTypeInfo* type_info;
  ort_type_info_provider->GetTypeInfo(&type_info);

  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->ModelAddOutput(ort_model_.get(), name, type_info), ort_api);
  return S_OK;
}


STDMETHODIMP OnnruntimeModel::JoinModel(_In_ IModel* other_model,
                                        _In_ const char* const* output_names,
                                        _In_ const char* const* input_names,
                                        size_t num_linkages,
                                        bool promote_unlinked_outputs,
                                        _In_ const char* const join_node_prefix) {
  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();
  auto ort_api = engine_factory_->UseOrtApi();
  
  RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->JoinModels(ort_model_.get(),
                                                       static_cast<OnnruntimeModel*>(other_model)->ort_model_.get(),
                                                       output_names,
                                                       input_names,
                                                       num_linkages,
                                                       promote_unlinked_outputs,
                                                       join_node_prefix),
                          ort_api);
  // reset the info so that it is recreated with the new information lazily
  info_ = nullptr;
  return S_OK;
}