Brianma/breaks (#2469)

* fix some more breaks

* learning model doesn't need lotusEnvironment and CPU shouldn't include dmlEP headers

* move dml checks out of winml and into the adapter

* better error handling

winml/lib/Api.Core/WinMLAdapter.cpp

@@ -57,8 +57,8 @@ class AbiSafeTensor : public Microsoft::WRL::RuntimeClass<
                           Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::ClassicCom>,
                           ITensor> {
  private:
-  onnxruntime::Tensor& tensor_;  // weak ref
-  ComPtr<IOrtValue> value_;      // strong ref
+  onnxruntime::Tensor& tensor_;              // weak ref
+  Microsoft::WRL::ComPtr<IOrtValue> 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<void*>(ml_value->GetMutable<std::map<std::string, int64_t>>());
       } else if (value_kind == TensorKind::Float) {
@@ -630,68 +629,80 @@ class WinMLAdapter : public Microsoft::WRL::RuntimeClass<
 #ifdef USE_DML
     auto impl = wil::MakeOrThrow<AbiCustomRegistryImpl>();
     *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<IMLOperatorRegistry> 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<int64_t>* 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<AbiSafeOrtValue>();
-  ort_value->get()->Init(tensor,
-                         onnxruntime::DataTypeImpl::GetType<onnxruntime::Tensor>(),
-                         onnxruntime::DataTypeImpl::GetType<onnxruntime::Tensor>()->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<int64_t>* 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<AbiSafeOrtValue>();
+    ort_value->get()->Init(tensor,
+                           onnxruntime::DataTypeImpl::GetType<onnxruntime::Tensor>(),
+                           onnxruntime::DataTypeImpl::GetType<onnxruntime::Tensor>()->GetDeleteFunc());
+
+    *gpu_value = ort_value.Detach();
+    return S_OK;
+#else
+    return E_NOTIMPL;
+#endif USE_DML
+  }
 
   HRESULT STDMETHODCALLTYPE CreateCPUMLValue(
       std::vector<int64_t>* 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<onnxruntime::IOBinding> binding_;
   std::vector<IOrtValue*> outputs_weak_;
-  std::vector<ComPtr<IOrtValue>> outputs_;
+  std::vector<Microsoft::WRL::ComPtr<IOrtValue>> 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;
 }

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
+}  // namespace Windows::AI::MachineLearning::Adapter
+
+#endif USE_DML

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;

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<ILearningModelOperatorProviderNative>();
 
-  // Retrieve the "operator abi" registry.
-  winrt::com_ptr<IMLOperatorRegistry> 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(), &registry));
+  return registry;
 }
 
 wfc::IVectorView<winml::ILearningModelFeatureDescriptor>