diff --git a/winml/lib/Api.Core/WinMLAdapter.cpp b/winml/lib/Api.Core/WinMLAdapter.cpp index 99d10f4318..b0b091ff5d 100644 --- a/winml/lib/Api.Core/WinMLAdapter.cpp +++ b/winml/lib/Api.Core/WinMLAdapter.cpp @@ -57,8 +57,8 @@ class AbiSafeTensor : public Microsoft::WRL::RuntimeClass< Microsoft::WRL::RuntimeClassFlags, ITensor> { private: - onnxruntime::Tensor& tensor_; // weak ref - ComPtr value_; // strong ref + onnxruntime::Tensor& tensor_; // weak ref + Microsoft::WRL::ComPtr value_; // strong ref public: AbiSafeTensor(onnxruntime::Tensor* tensor, @@ -589,8 +589,7 @@ class WinMLAdapter : public Microsoft::WRL::RuntimeClass< } else { THROW_HR(E_FAIL); } - } - else if (key_kind == TensorKind::String) { + } else if (key_kind == TensorKind::String) { if (value_kind == TensorKind::Int64) { return static_cast(ml_value->GetMutable>()); } else if (value_kind == TensorKind::Float) { @@ -630,68 +629,80 @@ class WinMLAdapter : public Microsoft::WRL::RuntimeClass< #ifdef USE_DML auto impl = wil::MakeOrThrow(); *registry = impl.Detach(); - return S_OK; + return S_OK; #else - return E_NOTIMPL; + return E_NOTIMPL; #endif USE_DML -} + } -void* STDMETHODCALLTYPE CreateGPUAllocationFromD3DResource(ID3D12Resource* pResource) override { + HRESULT STDMETHODCALLTYPE GetOperatorRegistry(ILearningModelOperatorProviderNative* operator_provider_native, IMLOperatorRegistry** registry) override { #ifdef USE_DML - return Dml::CreateGPUAllocationFromD3DResource(pResource); + // Retrieve the "operator abi" registry. + winrt::com_ptr operator_registry; + THROW_IF_FAILED(operator_provider_native->GetRegistry(operator_registry.put())); + *registry = operator_registry.detach(); + return S_OK; #else - return nullptr; + return E_NOTIMPL; #endif USE_DML -} + } -void STDMETHODCALLTYPE FreeGPUAllocation(void* ptr) override { + void* STDMETHODCALLTYPE CreateGPUAllocationFromD3DResource(ID3D12Resource* pResource) override { #ifdef USE_DML - Dml::FreeGPUAllocation(ptr); -#endif USE_DML -} - -HRESULT STDMETHODCALLTYPE CopyTensor( - onnxruntime::IExecutionProvider* provider, - ITensor* src, - ITensor* dst) override { -#ifdef USE_DML - ORT_THROW_IF_ERROR(Dml::CopyTensor(provider, src->get(), *(dst->getMutable()))); - return S_OK; + return Dml::CreateGPUAllocationFromD3DResource(pResource); #else - return E_NOTIMPL; + return nullptr; #endif USE_DML -} + } -HRESULT STDMETHODCALLTYPE CreateGPUMLValue( - void* execution_provider_allocated_resource, - onnxruntime::IExecutionProvider* provider, - std::vector* shape, - onnxruntime::MLDataType data_type, - IOrtValue** gpu_value) override { + void STDMETHODCALLTYPE FreeGPUAllocation(void* ptr) override { #ifdef USE_DML - THROW_HR_IF_MSG(WINML_ERR_INVALID_BINDING, - "DmlExecutionProvider" != provider->Type(), - "Cannot creat GPU tensor on CPU device"); - - onnxruntime::TensorShape tensor_shape(*shape); - - auto tensor = new onnxruntime::Tensor( - data_type, - tensor_shape, - execution_provider_allocated_resource, - provider->GetAllocator(0, ::OrtMemType::OrtMemTypeDefault)->Info()); - - auto ort_value = wil::MakeOrThrow(); - ort_value->get()->Init(tensor, - onnxruntime::DataTypeImpl::GetType(), - onnxruntime::DataTypeImpl::GetType()->GetDeleteFunc()); - - *gpu_value = ort_value.Detach(); - return S_OK; -#else - return E_NOTIMPL; + Dml::FreeGPUAllocation(ptr); #endif USE_DML -} + } + + HRESULT STDMETHODCALLTYPE CopyTensor( + onnxruntime::IExecutionProvider* provider, + ITensor* src, + ITensor* dst) override { +#ifdef USE_DML + ORT_THROW_IF_ERROR(Dml::CopyTensor(provider, src->get(), *(dst->getMutable()))); + return S_OK; +#else + return E_NOTIMPL; +#endif USE_DML + } + + HRESULT STDMETHODCALLTYPE CreateGPUMLValue( + void* execution_provider_allocated_resource, + onnxruntime::IExecutionProvider* provider, + std::vector* shape, + onnxruntime::MLDataType data_type, + IOrtValue** gpu_value) override { +#ifdef USE_DML + THROW_HR_IF_MSG(WINML_ERR_INVALID_BINDING, + "DmlExecutionProvider" != provider->Type(), + "Cannot creat GPU tensor on CPU device"); + + onnxruntime::TensorShape tensor_shape(*shape); + + auto tensor = new onnxruntime::Tensor( + data_type, + tensor_shape, + execution_provider_allocated_resource, + provider->GetAllocator(0, ::OrtMemType::OrtMemTypeDefault)->Info()); + + auto ort_value = wil::MakeOrThrow(); + ort_value->get()->Init(tensor, + onnxruntime::DataTypeImpl::GetType(), + onnxruntime::DataTypeImpl::GetType()->GetDeleteFunc()); + + *gpu_value = ort_value.Detach(); + return S_OK; +#else + return E_NOTIMPL; +#endif USE_DML + } HRESULT STDMETHODCALLTYPE CreateCPUMLValue( std::vector* shape, @@ -760,22 +771,22 @@ HRESULT STDMETHODCALLTYPE CreateGPUMLValue( return S_OK; } -// Override select shape inference functions which are incomplete in ONNX with versions that are complete, -// and are also used in DML kernel registrations. Doing this avoids kernel and shader creation being -// deferred until first evaluation. It also prevents a situation where inference functions in externally -// registered schema are reachable only after upstream schema have been revised in a later OS release, -// which would be a compatibility risk. -HRESULT STDMETHODCALLTYPE OverrideSchemaInferenceFunctions() override { + // Override select shape inference functions which are incomplete in ONNX with versions that are complete, + // and are also used in DML kernel registrations. Doing this avoids kernel and shader creation being + // deferred until first evaluation. It also prevents a situation where inference functions in externally + // registered schema are reachable only after upstream schema have been revised in a later OS release, + // which would be a compatibility risk. + HRESULT STDMETHODCALLTYPE OverrideSchemaInferenceFunctions() override { #ifdef USE_DML - static std::once_flag schema_override_once_flag; - std::call_once(schema_override_once_flag, []() { - SchemaInferenceOverrider::OverrideSchemaInferenceFunctions(); - }); - return S_OK; + static std::once_flag schema_override_once_flag; + std::call_once(schema_override_once_flag, []() { + SchemaInferenceOverrider::OverrideSchemaInferenceFunctions(); + }); + return S_OK; #else - return E_NOTIMPL; + return S_OK; // needs to return S_OK otherwise everything breaks because this gets called from the learningmodel constructor #endif USE_DML -} + } }; // namespace Windows::AI::MachineLearning::Adapter @@ -793,7 +804,7 @@ class IOBinding : public Microsoft::WRL::RuntimeClass< private: std::shared_ptr binding_; std::vector outputs_weak_; - std::vector> outputs_; + std::vector> outputs_; public: IOBinding(onnxruntime::IOBinding* binding) : binding_(binding) { @@ -883,12 +894,14 @@ InferenceSession::RegisterCustomRegistry( IMLOperatorRegistry* registry) { RETURN_HR_IF(S_OK, registry == nullptr); +#ifdef USE_DML auto custom_registries = GetLotusCustomRegistries(registry); // Register for (auto& custom_registry : custom_registries) { ORT_THROW_IF_ERROR(session_->RegisterCustomRegistry(custom_registry)); } +#endif USE_DML return S_OK; } diff --git a/winml/lib/Api.Core/inc/CustomRegistryHelper.h b/winml/lib/Api.Core/inc/CustomRegistryHelper.h index 4264f54017..de2987e676 100644 --- a/winml/lib/Api.Core/inc/CustomRegistryHelper.h +++ b/winml/lib/Api.Core/inc/CustomRegistryHelper.h @@ -3,6 +3,7 @@ #pragma once +#ifdef USE_DML #include "core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.h" namespace Windows::AI::MachineLearning::Adapter { @@ -24,4 +25,6 @@ GetLotusCustomRegistries( return {}; } -} // namespace Windows::AI::MachineLearning::Adapter \ No newline at end of file +} // namespace Windows::AI::MachineLearning::Adapter + +#endif USE_DML diff --git a/winml/lib/Api.Core/inc/WinMLAdapter.h b/winml/lib/Api.Core/inc/WinMLAdapter.h index f54165b56e..0a3d56b9f7 100644 --- a/winml/lib/Api.Core/inc/WinMLAdapter.h +++ b/winml/lib/Api.Core/inc/WinMLAdapter.h @@ -128,6 +128,7 @@ MIDL_INTERFACE("b19385e7-d9af-441a-ba7f-3993c7b1c9db") IWinMLAdapter : IUnknown // custom ops virtual HRESULT STDMETHODCALLTYPE GetCustomRegistry(IMLOperatorRegistry** registry) = 0; + virtual HRESULT STDMETHODCALLTYPE GetOperatorRegistry(ILearningModelOperatorProviderNative * operator_provider_native, IMLOperatorRegistry * *registry) = 0; // dml ep hooks virtual void* STDMETHODCALLTYPE CreateGPUAllocationFromD3DResource(ID3D12Resource* pResource) = 0; diff --git a/winml/lib/Api/LearningModel.cpp b/winml/lib/Api/LearningModel.cpp index 7636d33fd6..5223ed5c56 100644 --- a/winml/lib/Api/LearningModel.cpp +++ b/winml/lib/Api/LearningModel.cpp @@ -5,11 +5,8 @@ #include "LearningModel.h" -#include "core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.h" #include "TelemetryEvent.h" -#include "LotusEnvironment.h" - #include "MapFeatureDescriptor.h" #include "SequenceFeatureDescriptor.h" #include "TensorFeatureDescriptor.h" @@ -18,7 +15,7 @@ namespace winrt::Windows::AI::MachineLearning::implementation { LearningModel::LearningModel( const hstring& path, const winml::ILearningModelOperatorProvider op_provider) try : LearningModel(WinML::Strings::UTF8FromHString(path), - op_provider) { + op_provider) { } WINML_CATCH_ALL @@ -57,7 +54,7 @@ LearningModel::LearningModel( WINML_CATCH_ALL void LearningModel::Initialize() { - WINML_THROW_IF_FAILED(adapter_->CreateModelInfo(model_proto_.get(), model_info_.put())); + WINML_THROW_IF_FAILED(adapter_->CreateModelInfo(model_proto_.get(), model_info_.put())); } void LearningModel::LogCreationEvent(bool fromStream) { @@ -165,11 +162,9 @@ LearningModel::GetOperatorRegistry() { auto operator_provider_native = operator_provider_.as(); - // Retrieve the "operator abi" registry. - winrt::com_ptr operator_registry; - operator_provider_native->GetRegistry(operator_registry.put()); - - return operator_registry.get(); + IMLOperatorRegistry* registry = nullptr; + WINML_THROW_IF_FAILED(adapter_->GetOperatorRegistry(operator_provider_native.get(), ®istry)); + return registry; } wfc::IVectorView