Add Named Dimension Override API to LearningModelSessionOptions (WinML) (#4606)

Co-authored-by: Ori Levari <orlevari@microsoft.com>

winml/api/Windows.AI.MachineLearning.idl

@@ -33,7 +33,7 @@ import "windows.storage.idl";
 
 namespace ROOT_NS.AI.MachineLearning 
 {
-    [contractversion(3)]
+    [contractversion(4)]
     apicontract MachineLearningContract{};
 
     //! Forward declarations
@@ -213,6 +213,15 @@ namespace ROOT_NS.AI.MachineLearning
         //! BatchSizeOverride > 0 indicates the size of batch that will be used to override the model batch size and optimize evaluations.
         UInt32 BatchSizeOverride { get; set; };
 
+        [contract(MachineLearningContract, 4)]
+        {
+            //! The OverrideNamedDimension method will allow the model compiler to use constant batch size performance optimizations when setting up the LearningModelSession.
+            //! The caller can specify the size of the dimension for a given named dimension.
+            //! dimension = 0 indicates that the dimension present in the model should be honored.
+            //! dimension > 0 indicates the size of the dimension that will be used to override the model "name" dimension and optimize evaluations.
+            void OverrideNamedDimension(String name, UInt32 dimension);
+        }
+
         [contract(MachineLearningContract, 3)]
         {
             //! The CloseModelOnSessionCreation option will allow the LearningModelSession to take ownership of the LearningModel's

winml/lib/Api.Ort/OnnxruntimeEngineBuilder.cpp

@@ -41,6 +41,12 @@ STDMETHODIMP OnnxruntimeEngineBuilder::CreateEngine(_winml::IEngine** out) {
     RETURN_HR_IF_NOT_OK_MSG(ort_api->AddFreeDimensionOverride(session_options.get(), DATA_BATCH, batch_size_override_.value()),
                             ort_api);
   }
+  if (named_dimension_overrides_) {
+    for (const auto& override : named_dimension_overrides_) {
+      std::string narrow_name = std::wstring_convert<std::codecvt_utf8<wchar_t>>().to_bytes(override.Key().c_str());
+      ort_api->AddFreeDimensionOverrideByName(session_options.get(), narrow_name.c_str(), override.Value());
+    }
+  }
 
   RETURN_HR_IF_NOT_OK_MSG(ort_api->SetIntraOpNumThreads(session_options.get(), intra_op_num_threads_override_), ort_api);
 
@@ -81,6 +87,12 @@ STDMETHODIMP OnnxruntimeEngineBuilder::SetBatchSizeOverride(uint32_t batch_size_
   return S_OK;
 }
 
+
+STDMETHODIMP OnnxruntimeEngineBuilder::SetNamedDimensionOverrides(wfc::IMapView<winrt::hstring, uint32_t> named_dimension_overrides) {
+  named_dimension_overrides_ = std::move(named_dimension_overrides);
+  return S_OK;
+}
+  
 STDMETHODIMP OnnxruntimeEngineBuilder::SetIntraOpNumThreadsOverride(uint32_t intra_op_num_threads) {
   intra_op_num_threads_override_ = intra_op_num_threads;
   return S_OK;

winml/lib/Api.Ort/OnnxruntimeEngineBuilder.h

@@ -26,6 +26,9 @@ class OnnxruntimeEngineBuilder : public Microsoft::WRL::RuntimeClass<
   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);
 
@@ -38,6 +41,7 @@ class OnnxruntimeEngineBuilder : public Microsoft::WRL::RuntimeClass<
   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_;
 };
 

winml/lib/Api/LearningModelSession.cpp

@@ -110,6 +110,7 @@ void LearningModelSession::Initialize() {
     WINML_THROW_IF_FAILED(engine_builder->SetMetacommandsEnabled(device_impl->MetacommandsEnabled()));
   }
 
+
   // Make onnxruntime apply the batch size override, if any
   if (session_options_) {
     if (session_options_.BatchSizeOverride() != 0) {
@@ -119,6 +120,12 @@ void LearningModelSession::Initialize() {
     uint32_t numIntraOpThreads = session_options_.as<WINMLP::LearningModelSessionOptions>()->GetIntraOpNumThreads();
     WINML_THROW_IF_FAILED(engine_builder->SetIntraOpNumThreadsOverride(numIntraOpThreads)
     );
+    
+    // Make onnxruntime apply named dimension overrides, if any
+    com_ptr<winmlp::LearningModelSessionOptions> session_options_impl = session_options_.as<winmlp::LearningModelSessionOptions>();
+    if (session_options_impl && session_options_impl->NamedDimensionOverrides().Size() > 0) {
+      WINML_THROW_IF_FAILED(engine_builder->SetNamedDimensionOverrides(session_options_impl->NamedDimensionOverrides()));
+    }
   } else {
     // Onnxruntime will use half the number of concurrent threads supported on the system
     // by default. This causes MLAS to not exercise every logical core.

winml/lib/Api/LearningModelSessionOptions.cpp

@@ -24,6 +24,14 @@ void LearningModelSessionOptions::CloseModelOnSessionCreation(bool value) {
   close_model_on_session_creation_ = value;
 }
 
+wfc::IMapView<winrt::hstring, uint32_t> LearningModelSessionOptions::NamedDimensionOverrides() {
+  return named_dim_overrides_.GetView();
+}
+
+void LearningModelSessionOptions::OverrideNamedDimension(winrt::hstring name, uint32_t value) {
+  named_dim_overrides_.Insert(name, value);
+}
+
 uint32_t LearningModelSessionOptions::GetIntraOpNumThreads() {
   return intra_op_num_threads_override_;
 }
@@ -32,4 +40,4 @@ STDMETHODIMP LearningModelSessionOptions::SetIntraOpNumThreadsOverride(uint32_t
   intra_op_num_threads_override_ = intraOpNumThreads;
   return S_OK;
 }
-}  // namespace WINMLP
+}  // namespace WINMLP

winml/lib/Api/LearningModelSessionOptions.h

@@ -17,6 +17,9 @@ struct LearningModelSessionOptions : LearningModelSessionOptionsT<LearningModelS
 
   bool CloseModelOnSessionCreation();
   void CloseModelOnSessionCreation(bool value);
+  
+  wfc::IMapView<winrt::hstring, uint32_t> NamedDimensionOverrides();
+  void OverrideNamedDimension(winrt::hstring name, uint32_t value);
 
   STDMETHOD(SetIntraOpNumThreadsOverride)
   (uint32_t intraOpNumThreads);
@@ -47,6 +50,13 @@ struct LearningModelSessionOptions : LearningModelSessionOptionsT<LearningModelS
   // The default value here is False so that models are not automatically closed on session creation.
   bool close_model_on_session_creation_ = false;
 
+  // Map of named input dimensions to concrete values. 
+  // This informs the engine when the developer wants to explictily set a named dimension to a fixed value.
+
+  // 0    : the dimension present in the model should be honored.
+  // 1...n: override the named input dimension to the given value and optimize evaluations.
+  wfc::IMap<winrt::hstring, uint32_t> named_dim_overrides_ = winrt::single_threaded_map<winrt::hstring, uint32_t>();
+
   // The intra operator num threads property is used to control the number of threads used in the threadpool for intra operator calculations.
   // The default value here is the maximum number of logical cores to ensure that the default behavior of WinML always runs the fastest.
   // WARNING: Setting a number higher than the maximum number of logical cores may result in an inefficient threadpool

winml/lib/Common/inc/iengine.h

@@ -168,6 +168,9 @@ IEngineBuilder : IUnknown {
   STDMETHOD(SetBatchSizeOverride)
   (uint32_t batch_size_override) PURE;
 
+  STDMETHOD(SetNamedDimensionOverrides)
+  (wfc::IMapView<winrt::hstring, uint32_t> named_dimension_overrides) PURE;
+
   STDMETHOD(SetIntraOpNumThreadsOverride)
   (uint32_t intra_op_num_threads) PURE;
 

winml/test/api/LearningModelSessionAPITest.cpp

@@ -311,6 +311,40 @@ static void EvaluateSessionAndCloseModel()
     WINML_EXPECT_NO_THROW(::EvaluateSessionAndCloseModelHelper(LearningModelDeviceKind::Cpu, false));
 }
 
+static void NamedDimensionOverride() 
+{
+  LearningModel model = nullptr;
+  WINML_EXPECT_NO_THROW(APITest::LoadModel(L"fns-candy.onnx", model));
+
+  LearningModelDevice device(nullptr);
+  WINML_EXPECT_NO_THROW(device = LearningModelDevice(LearningModelDeviceKind::Cpu));
+
+  // the model input shape. the batch size, n, is overriden to 5
+  int64_t n = 5, c = 3, h = 720, w = 720;
+
+  LearningModelSessionOptions options;
+  options.OverrideNamedDimension(L"None", n);
+  
+  // Verifies that if a Dim name doesn't exist the named dimension override does nothing
+  options.OverrideNamedDimension(L"DimNameThatDoesntExist", n);
+
+  LearningModelSession session(nullptr);
+  WINML_EXPECT_NO_THROW(session = LearningModelSession(model, device, options));
+
+  ILearningModelFeatureDescriptor descriptor = model.InputFeatures().GetAt(0);
+  TensorFeatureDescriptor tensorDescriptor = nullptr;
+  descriptor.as(tensorDescriptor);
+  std::vector<int64_t> shape{n,c,h,w};
+  int64_t size = n*c*h*w;
+  std::vector<float> buffer;
+  buffer.resize(static_cast<size_t>(size));
+  auto featureValue = TensorFloat::CreateFromIterable(shape, winrt::single_threaded_vector<float>(std::move(buffer)));
+  LearningModelBinding binding(session);
+  binding.Bind(descriptor.Name(), featureValue);
+
+  WINML_EXPECT_NO_THROW(session.Evaluate(binding, L""));
+}
+
 static void CloseSession()
 {
     LearningModel learningModel = nullptr;
@@ -435,6 +469,7 @@ const LearningModelSessionAPITestsApi& getapi() {
     CreateSessionWithCastToFloat16InModel,
     CreateSessionWithFloat16InitializersInModel,
     EvaluateSessionAndCloseModel,
+    NamedDimensionOverride,
     CloseSession,
     SetIntraOpNumThreads
   };

winml/test/api/LearningModelSessionAPITest.h

@@ -18,6 +18,7 @@ struct LearningModelSessionAPITestsApi {
   VoidTest CreateSessionWithCastToFloat16InModel;
   VoidTest CreateSessionWithFloat16InitializersInModel;
   VoidTest EvaluateSessionAndCloseModel;
+  VoidTest OverrideNamedDimension;
   VoidTest CloseSession;
   VoidTest SetIntraOpNumThreads;
 };
@@ -39,5 +40,7 @@ WINML_TEST(LearningModelSessionAPITests, CreateSessionDeviceDirectXMinimumPower)
 WINML_TEST(LearningModelSessionAPITests, CreateSessionWithCastToFloat16InModel)
 WINML_TEST(LearningModelSessionAPITests, CreateSessionWithFloat16InitializersInModel)
 WINML_TEST(LearningModelSessionAPITests, AdapterIdAndDevice)
+WINML_TEST(LearningModelSessionAPITests, OverrideNamedDimension)
+WINML_TEST(LearningModelSessionAPITests, CloseSession)
 WINML_TEST(LearningModelSessionAPITests, SetIntraOpNumThreads)
 WINML_TEST_CLASS_END()