onnxruntime/winml/lib/Api.Ort/OnnxruntimeEngineBuilder.h
Ori Levari e6ef3653a7
Add Named Dimension Override API to LearningModelSessionOptions (WinML) (#4606)
Co-authored-by: Ori Levari <orlevari@microsoft.com>
2020-08-03 16:04:21 -07:00

48 lines
No EOL
1.4 KiB
C++

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "iengine.h"
namespace _winml {
class OnnxruntimeEngineBuilder : public Microsoft::WRL::RuntimeClass<
Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::ClassicCom>,
IEngineBuilder> {
public:
HRESULT RuntimeClassInitialize(_In_ OnnxruntimeEngineFactory* engine);
STDMETHOD(SetD3D12Resources)
(ID3D12Device* device, ID3D12CommandQueue* queue);
STDMETHOD(SetMetacommandsEnabled)
(int enabled);
STDMETHOD(GetD3D12Device)
(_Outptr_ ID3D12Device** device);
STDMETHOD(GetID3D12CommandQueue)
(_Outptr_ ID3D12CommandQueue** queue);
STDMETHOD(SetBatchSizeOverride)
(uint32_t batch_size_override);
STDMETHOD(SetNamedDimensionOverrides)
(wfc::IMapView<winrt::hstring, uint32_t> named_dimension_overrides);
STDMETHOD(SetIntraOpNumThreadsOverride)
(uint32_t intra_op_num_threads);
STDMETHOD(CreateEngine)
(_Outptr_ IEngine** out);
private:
Microsoft::WRL::ComPtr<OnnxruntimeEngineFactory> engine_factory_;
Microsoft::WRL::ComPtr<ID3D12Device> device_ = nullptr;
Microsoft::WRL::ComPtr<ID3D12CommandQueue> queue_ = nullptr;
bool metacommands_enabled_ = true;
std::optional<uint32_t> batch_size_override_;
wfc::IMapView<winrt::hstring, uint32_t> named_dimension_overrides_;
uint32_t intra_op_num_threads_override_;
};
} // namespace _winml