From 1a20b6d13ec3bfd9e88ea70c955db023095f4bd8 Mon Sep 17 00:00:00 2001
From: Ryan Hill <ryanhill@microsoft.com>
Date: Wed, 12 May 2021 17:17:32 -0700
Subject: [PATCH] Add in a couple more missing GetType methods

---
 onnxruntime/core/framework/provider_bridge_ort.cc             | 2 ++
 .../core/providers/shared_library/provider_bridge_provider.cc | 4 ++++
 .../core/providers/shared_library/provider_interfaces.h       | 2 ++
 3 files changed, 8 insertions(+)

diff --git a/onnxruntime/core/framework/provider_bridge_ort.cc b/onnxruntime/core/framework/provider_bridge_ort.cc
index 853fc62df8..e0dec395f6 100644
--- a/onnxruntime/core/framework/provider_bridge_ort.cc
+++ b/onnxruntime/core/framework/provider_bridge_ort.cc
@@ -492,6 +492,8 @@ struct ProviderHostImpl : ProviderHost {
   MLDataType DataTypeImpl__GetType_uint64() override { return DataTypeImpl::GetType<uint64_t>(); }
   MLDataType DataTypeImpl__GetType_float() override { return DataTypeImpl::GetType<float>(); }
   MLDataType DataTypeImpl__GetType_double() override { return DataTypeImpl::GetType<double>(); }
+  MLDataType DataTypeImpl__GetType_BFloat16() override { return DataTypeImpl::GetType<BFloat16>(); }
+  MLDataType DataTypeImpl__GetType_MLFloat16() override { return DataTypeImpl::GetType<MLFloat16>(); }
   MLDataType DataTypeImpl__GetTensorType_bool() override { return DataTypeImpl::GetTensorType<bool>(); }
   MLDataType DataTypeImpl__GetTensorType_int8() override { return DataTypeImpl::GetTensorType<int8_t>(); }
   MLDataType DataTypeImpl__GetTensorType_uint8() override { return DataTypeImpl::GetTensorType<uint8_t>(); }
diff --git a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc
index 513b950110..4afab050f4 100644
--- a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc
+++ b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc
@@ -91,6 +91,10 @@ MLDataType DataTypeImpl::GetType<float>() { return Provider_GetHost()->DataTypeI
 template <>
 MLDataType DataTypeImpl::GetType<double>() { return Provider_GetHost()->DataTypeImpl__GetType_double(); }
 template <>
+MLDataType DataTypeImpl::GetType<BFloat16>() { return Provider_GetHost()->DataTypeImpl__GetType_BFloat16(); }
+template <>
+MLDataType DataTypeImpl::GetType<MLFloat16>() { return Provider_GetHost()->DataTypeImpl__GetType_MLFloat16(); }
+template <>
 MLDataType DataTypeImpl::GetTensorType<bool>() { return Provider_GetHost()->DataTypeImpl__GetTensorType_bool(); }
 template <>
 MLDataType DataTypeImpl::GetTensorType<int8_t>() { return Provider_GetHost()->DataTypeImpl__GetTensorType_int8(); }
diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h
index 01be5b2a30..ab8b5dba39 100644
--- a/onnxruntime/core/providers/shared_library/provider_interfaces.h
+++ b/onnxruntime/core/providers/shared_library/provider_interfaces.h
@@ -413,6 +413,8 @@ struct ProviderHost {
   virtual MLDataType DataTypeImpl__GetType_uint64() = 0;
   virtual MLDataType DataTypeImpl__GetType_float() = 0;
   virtual MLDataType DataTypeImpl__GetType_double() = 0;
+  virtual MLDataType DataTypeImpl__GetType_BFloat16() = 0;
+  virtual MLDataType DataTypeImpl__GetType_MLFloat16() = 0;
   virtual MLDataType DataTypeImpl__GetTensorType_bool() = 0;
   virtual MLDataType DataTypeImpl__GetTensorType_int8() = 0;
   virtual MLDataType DataTypeImpl__GetTensorType_uint8() = 0;