Add ability for pre-packed weights of shared initializers to be shared across sessions

csharp/src/Microsoft.ML.OnnxRuntime/InferenceSession.cs

@@ -47,6 +47,21 @@ public InferenceSession(string modelPath)
             Init(modelPath, _builtInSessionOptions);
         }
 
+        /// <summary>
+        /// Constructs an InferenceSession from a model file and it will use 
+        /// the provided pre-packed weights container to store and share pre-packed buffers 
+        /// of shared initializers across sessions if any.
+        /// </summary>
+        /// <param name="modelPath">Model path</param>
+        /// <param name="prepackedWeightsContainer">Instance of PrepackedWeightsContainer. 
+        /// Lifetime of 'prepackedWeightsContainer' must be
+        /// managed by the user and it must outlive any sessions reliant on it</param>
+        public InferenceSession(string modelPath, PrePackedWeightsContainer prepackedWeightsContainer)
+        {
+            _builtInSessionOptions = new SessionOptions(); // need to be disposed
+            Init(modelPath, _builtInSessionOptions, prepackedWeightsContainer);
+        }
+
 
         /// <summary>
         /// Constructs an InferenceSession from a model file, with some additional session options
@@ -58,6 +73,23 @@ public InferenceSession(string modelPath, SessionOptions options)
             Init(modelPath, options);
         }
 
+
+        /// <summary>
+        /// Constructs an InferenceSession from a model file, with some additional session options
+        /// and it will use the provided pre-packed weights container to store and share pre-packed buffers 
+        /// of shared initializers across sessions if any.
+        /// </summary>
+        /// <param name="modelPath">Model path</param>
+        /// <param name="options">Session options</param>
+        /// <param name="prepackedWeightsContainer">Instance of PrepackedWeightsContainer. 
+        /// Lifetime of 'prepackedWeightsContainer' must be
+        /// managed by the user and it must outlive any sessions reliant on it</param>
+        public InferenceSession(string modelPath, SessionOptions options,
+            PrePackedWeightsContainer prepackedWeightsContainer)
+        {
+            Init(modelPath, options, prepackedWeightsContainer);
+        }
+
         /// <summary>
         /// Constructs an InferenceSession from a model data in byte array
         /// </summary>
@@ -68,6 +100,21 @@ public InferenceSession(byte[] model)
             Init(model, _builtInSessionOptions);
         }
 
+        /// <summary>
+        /// Constructs an InferenceSession from a model data (in byte array) and it will use 
+        /// the provided pre-packed weights container to store and share pre-packed buffers 
+        /// of shared initializers across sessions if any.
+        /// </summary>
+        /// <param name="model">Model as byte array</param>
+        /// <param name="prepackedWeightsContainer">Instance of PrepackedWeightsContainer. 
+        /// Lifetime of 'prepackedWeightsContainer' must be
+        /// managed by the user and it must outlive any sessions reliant on it</param>
+        public InferenceSession(byte[] model, PrePackedWeightsContainer prepackedWeightsContainer)
+        {
+            _builtInSessionOptions = new SessionOptions(); // need to be disposed
+            Init(model, _builtInSessionOptions, prepackedWeightsContainer);
+        }
+
         /// <summary>
         /// Constructs an InferenceSession from a model data in byte array, with some additional session options
         /// </summary>
@@ -78,6 +125,21 @@ public InferenceSession(byte[] model, SessionOptions options)
             Init(model, options);
         }
 
+        /// <summary>
+        /// Constructs an InferenceSession from a model data (in byte array) with some additional
+        /// session options and it will use the provided pre-packed weights container to store
+        /// and share pre-packed buffers of shared initializers across sessions if any.
+        /// </summary>
+        /// <param name="model">Model as byte array</param>
+        /// <param name="options">Session Options</param>
+        /// <param name="prepackedWeightsContainer">Instance of PrepackedWeightsContainer. 
+        /// Lifetime of 'prepackedWeightsContainer' must be
+        /// managed by the user and it must outlive any sessions reliant on it</param>
+        public InferenceSession(byte[] model, SessionOptions options,
+                                PrePackedWeightsContainer prepackedWeightsContainer)
+        {
+            Init(model, options, prepackedWeightsContainer);
+        }
         /// <summary>
         /// Meta data regarding the input nodes, keyed by input names
         /// </summary>
@@ -310,7 +372,7 @@ public void Run(
             IReadOnlyCollection<NamedOnnxValue> outputs,
             RunOptions options)
         {
-            using(var cleanupList = new DisposableList<IDisposable>())
+            using (var cleanupList = new DisposableList<IDisposable>())
             {
                 var inputNamesArray = ConvertNamesToUtf8(inputs, i => i.Name, cleanupList);
                 var inputValuesArray = GetOrtValuesHandles(inputs, cleanupList);
@@ -367,7 +429,7 @@ public void Run(
                 throw new ArgumentException($"Length of {nameof(outputNames)} ({outputNames.Count}) must match that of {nameof(outputValues)} ({outputValues.Count}).");
             }
 
-            using(var cleanupList = new DisposableList<IDisposable>())
+            using (var cleanupList = new DisposableList<IDisposable>())
             {
                 // prepare inputs
                 var inputNamesArray = ConvertNamesToUtf8(inputs, i => i.Name, cleanupList);
@@ -428,7 +490,7 @@ public void Run(
                 throw new ArgumentException($"Length of {nameof(inputNames)} ({inputNames.Count}) must match that of {nameof(inputValues)} ({inputValues.Count}).");
             }
 
-            using(var cleanupList = new DisposableList<IDisposable>())
+            using (var cleanupList = new DisposableList<IDisposable>())
             {
                 // prepare inputs
                 var inputNamesArray = ConvertNamesToUtf8(inputNames, n => n, cleanupList);
@@ -515,7 +577,8 @@ public IDisposableReadOnlyCollection<DisposableNamedOnnxValue> RunWithBindingAnd
                         var ortValue = ortValues.ElementAt(i);
                         result.Add(DisposableNamedOnnxValue.CreateFromOrtValue(outputNames[i], ortValue));
                     }
-                } catch(Exception e)
+                }
+                catch (Exception e)
                 {
                     result.Dispose();
                     throw e;
@@ -535,7 +598,7 @@ public string EndProfiling()
             NativeApiStatus.VerifySuccess(NativeMethods.OrtSessionEndProfiling(_nativeHandle,
                                                                    allocator.Pointer,
                                                                    out nameHandle));
-            using(var allocation = new OrtMemoryAllocation(allocator, nameHandle, 0))
+            using (var allocation = new OrtMemoryAllocation(allocator, nameHandle, 0))
             {
                 return NativeOnnxValueHelper.StringFromNativeUtf8(nameHandle);
             }
@@ -609,8 +672,8 @@ private IntPtr[] GetOrtValuesHandles(IReadOnlyCollection<FixedBufferOnnxValue> v
         }
 
 
-         private DisposableList<OrtValue> RunImpl(RunOptions options, IntPtr[] inputNames, IntPtr[] inputValues, IntPtr[] outputNames,
-            DisposableList<IDisposable> cleanupList)
+        private DisposableList<OrtValue> RunImpl(RunOptions options, IntPtr[] inputNames, IntPtr[] inputValues, IntPtr[] outputNames,
+           DisposableList<IDisposable> cleanupList)
         {
             var ortValues = new DisposableList<OrtValue>(outputNames.Length);
             cleanupList.Add(ortValues);
@@ -680,31 +743,57 @@ public ModelMetadata ModelMetadata
         /// </summary>
         public ulong ProfilingStartTimeNs
         {
-           get
-           {
-               return _profilingStartTimeNs;
-           }
-        }  
+            get
+            {
+                return _profilingStartTimeNs;
+            }
+        }
 
         #endregion
 
         #region private methods
 
-        private void Init(string modelPath, SessionOptions options)
+        private void Init(string modelPath, SessionOptions options,
+                          PrePackedWeightsContainer prepackedWeightsContainer = null)
         {
             var envHandle = OrtEnv.Handle;
             var session = IntPtr.Zero;
-            NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSession(envHandle, NativeMethods.GetPlatformSerializedString(modelPath), options.Handle, out session));
+
+            if (prepackedWeightsContainer == null)
+            {
+                NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSession(envHandle, NativeMethods.GetPlatformSerializedString(modelPath),
+                    options.Handle, out session));
+            }
+
+            else
+            {
+                NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSessionWithPrepackedWeightsContainer(
+                    envHandle, NativeMethods.GetPlatformSerializedString(modelPath),
+                    options.Handle, prepackedWeightsContainer.Pointer, out session));
+            }
 
             InitWithSessionHandle(session, options);
         }
 
-        private void Init(byte[] modelData, SessionOptions options)
+        private void Init(byte[] modelData, SessionOptions options,
+                          PrePackedWeightsContainer prepackedWeightsContainer = null)
         {
             var envHandle = OrtEnv.Handle;
             var session = IntPtr.Zero;
 
-            NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSessionFromArray(envHandle, modelData, (UIntPtr)modelData.Length, options.Handle, out session));
+            if (prepackedWeightsContainer == null)
+            {
+
+                NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSessionFromArray(envHandle, modelData, (UIntPtr)modelData.Length, options.Handle, out session));
+            }
+
+            else
+            {
+                NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSessionFromArrayWithPrepackedWeightsContainer(
+                    envHandle, modelData, (UIntPtr)modelData.Length, options.Handle, prepackedWeightsContainer.Pointer,
+                    out session));
+
+            }
 
             InitWithSessionHandle(session, options);
         }
@@ -757,8 +846,8 @@ private void InitWithSessionHandle(IntPtr session, SessionOptions options)
                 // set profiling's start time
                 UIntPtr startTime = UIntPtr.Zero;
                 NativeApiStatus.VerifySuccess(NativeMethods.OrtSessionGetProfilingStartTimeNs(_nativeHandle,
-                                                                    out startTime)); 
-                _profilingStartTimeNs = (ulong) startTime;
+                                                                    out startTime));
+                _profilingStartTimeNs = (ulong)startTime;
             }
             catch (OnnxRuntimeException e)
             {
@@ -821,7 +910,7 @@ private string GetOverridableInitializerName(ulong index)
                                             (UIntPtr)index,
                                             allocator.Pointer,
                                             out nameHandle));
-            using(var ortAllocation = new OrtMemoryAllocation(allocator, nameHandle, 0))
+            using (var ortAllocation = new OrtMemoryAllocation(allocator, nameHandle, 0))
             {
                 str = NativeOnnxValueHelper.StringFromNativeUtf8(nameHandle);
             }
@@ -963,7 +1052,7 @@ public void Dispose()
         /// <param name="disposing">true if invoked from Dispose() method</param>
         protected virtual void Dispose(bool disposing)
         {
-            if(_disposed)
+            if (_disposed)
             {
                 return;
             }
@@ -1137,7 +1226,7 @@ internal ModelMetadata(InferenceSession session)
                     }
 
                     // Process each key via the stored key handles
-                    foreach(var allocation in ortAllocationKeys)
+                    foreach (var allocation in ortAllocationKeys)
                     {
                         IntPtr keyHandle = allocation.Pointer;
                         IntPtr valueHandle = IntPtr.Zero;
@@ -1160,9 +1249,9 @@ internal ModelMetadata(InferenceSession session)
             {
 
                 // Free ModelMetadata handle
-                 NativeMethods.OrtReleaseModelMetadata(modelMetadataHandle);
+                NativeMethods.OrtReleaseModelMetadata(modelMetadataHandle);
 
-             }
+            }
 
         }