Compile API: output model and initializer stream write functions

csharp/src/Microsoft.ML.OnnxRuntime/CompileModel.shared.cs

@@ -130,6 +130,287 @@ public void SetFlags(OrtCompileApiFlags flags)
                 NativeMethods.CompileApi.OrtModelCompilationOptions_SetFlags(handle, (uint)flags));
         }
 
+        /// <summary>
+        /// Sets information related to EP context binary file. The Ep uses this information to decide the
+        /// location and context binary file name when compiling with both the input and output models
+        /// stored in buffers.
+        /// </summary>
+        /// <param name="outputDirectory">Path to the model directory.</param>
+        /// <param name="modelName">The name of the model.</param>
+        public void SetEpContextBinaryInformation(string outputDirectory, string modelName)
+        {
+            var platformOutputDirectory = NativeOnnxValueHelper.GetPlatformSerializedString(outputDirectory);
+            var platformModelName = NativeOnnxValueHelper.GetPlatformSerializedString(modelName);
+            NativeApiStatus.VerifySuccess(
+                NativeMethods.CompileApi.OrtModelCompilationOptions_SetEpContextBinaryInformation(
+                    handle, platformOutputDirectory, platformModelName));
+        }
+
+        /// <summary>
+        /// Delegate to write/save a buffer containing ONNX model bytes to a custom destination.
+        /// </summary>
+        /// <param name="buffer">The buffer to write out.</param>
+        public delegate void WriteBufferDelegate(ReadOnlySpan<byte> buffer);
+
+        /// <summary>
+        /// Sets a delegate that is called by ORT to write out the output model's serialized ONNX bytes.
+        /// The provided delegate may be called repeatedly until the entire output model has been written out.
+        /// Each call to the delegate is expected to consume/handle the entire input buffer.
+        /// </summary>
+        /// <param name="writeBufferDelegate">The delegate called by ORT to write out the model.</param>
+        public void SetOutputModelWriteBufferDelegate(WriteBufferDelegate writeBufferDelegate)
+        {
+            _writeBufferDelegateState?.Dispose();
+            _writeBufferDelegateState =
+                new DelegateResources<WriteBufferConnector, NativeMethods.DOrtWriteBufferDelegate>(
+                    new WriteBufferConnector(writeBufferDelegate),
+                    new NativeMethods.DOrtWriteBufferDelegate(WriteBufferConnector.WriteBufferDelegateWrapper));
+
+            IntPtr funcPtr = _writeBufferDelegateState.GetFunctionPointerForDelegate();
+
+            NativeApiStatus.VerifySuccess(
+                NativeMethods.CompileApi.OrtModelCompilationOptions_SetOutputModelWriteFunc(
+                    handle,
+                    funcPtr,
+                    _writeBufferDelegateState.GetConnectorHandleAsPointer()));
+        }
+
+        /// <summary>
+        /// Delegate called by ORT for every initializer when generating the compiled model.
+        /// The delegate allows the user to determine whether the initializer should be stored within the compiled
+        /// model or externally in a file. If the delegate chooses to store an initializer externally, the delegate
+        /// implementation is responsible for writing the initializer data to a file.
+        /// </summary>
+        /// <param name="initializerName">The initializer's name.</param>
+        /// <param name="initializerValue">The readonly OrtValue instance containing the data, type, and
+        /// shape of the initializer.</param>
+        /// <param name="optionalExternalInfo">May be null. If the initializer is originally stored externally,
+        /// this contains the file path, file offset, and data size. Otherwise, this is null.</param>
+        /// <returns>A new OrtExternalInitializerInfo indicating the new location of the initializer.
+        /// Returns null if the initializer should be stored within the generated compiled model.</returns>
+        /// <remarks>The return value may be null.</remarks>
+        public delegate OrtExternalInitializerInfo HandleInitializerDelegate(
+            string initializerName,
+            IReadOnlyOrtValue initializerValue,
+            IReadOnlyExternalInitializerInfo optionalExternalInfo);
+
+        /// <summary>
+        /// Sets a delegate that is called by ORT for every initializer when generating the compiled model.
+        /// The delegate allows the user to determine whether the initializer should be stored within the compiled
+        /// model or externally in a file. If the delegate chooses to store an initializer externally, the delegate
+        /// implementation is responsible for writing the initializer data to a file.
+        /// </summary>
+        /// <param name="handleInitializerDelegate">The delegate called by ORT for every initializer.</param>
+        public void SetOutputModelHandleInitializerDelegate(HandleInitializerDelegate handleInitializerDelegate)
+        {
+            _handleInitializerDelegateState?.Dispose();
+            _handleInitializerDelegateState =
+                new DelegateResources<HandleInitializerConnector, NativeMethods.DOrtHandleInitializerDataDelegate>(
+                    new HandleInitializerConnector(handleInitializerDelegate),
+                    new NativeMethods.DOrtHandleInitializerDataDelegate(
+                        HandleInitializerConnector.HandleInitializerDelegateWrapper));
+
+            IntPtr funcPtr = _handleInitializerDelegateState.GetFunctionPointerForDelegate();
+
+            NativeApiStatus.VerifySuccess(
+                NativeMethods.CompileApi.OrtModelCompilationOptions_SetOutputModelHandleInitializerFunc(
+                    handle,
+                    funcPtr,
+                    _handleInitializerDelegateState.GetConnectorHandleAsPointer()));
+        }
+
+        #region Delegate helpers
+        /// <summary>
+        /// Class to bridge the C# and native worlds for the "write buffer" delegate
+        /// </summary>
+        private class WriteBufferConnector
+        {
+            private readonly WriteBufferDelegate _csharpDelegate;
+
+            internal WriteBufferConnector(WriteBufferDelegate writeBufferDelegate)
+            {
+                _csharpDelegate = writeBufferDelegate;
+            }
+
+            public static IntPtr WriteBufferDelegateWrapper(IntPtr /* void* */ state,
+                                                            IntPtr /* const void* */ buffer,
+                                                            UIntPtr /* size_t */ bufferNumBytes)
+            {
+                try
+                {
+
+                    WriteBufferConnector connector = (WriteBufferConnector)GCHandle.FromIntPtr(state).Target;
+                    ReadOnlySpan<byte> bufferSpan;
+
+                    unsafe
+                    {
+                        // NOTE: A Span<byte> can only view 2GB of data. This is fine because ORT does not write out
+                        // chunks that large. However, if we ever need to, the solution is to just write a loop here
+                        // that repeatedly calls the delegate with smaller chunks of data.
+                        bufferSpan = new ReadOnlySpan<byte>(buffer.ToPointer(), checked((int)bufferNumBytes));
+                    }
+
+                    connector._csharpDelegate(bufferSpan);
+                }
+                catch (Exception ex)
+                {
+                    var error = $"The C# WriteBuffer delegate threw an exception: {ex.Message}";
+                    IntPtr status = NativeMethods.OrtCreateStatus((uint)ErrorCode.Fail,
+                                            NativeOnnxValueHelper.StringToZeroTerminatedUtf8(error));
+                    return status;
+                }
+
+                return IntPtr.Zero;
+            }
+        }
+
+        /// <summary>
+        /// Class to bridge the C# and native worlds for the "handle initializer" delegate
+        /// </summary>
+        private class HandleInitializerConnector
+        {
+            private readonly HandleInitializerDelegate _csharpDelegate;
+
+            internal HandleInitializerConnector(HandleInitializerDelegate handleInitializerDelegate)
+            {
+                _csharpDelegate = handleInitializerDelegate;
+            }
+
+            public static IntPtr HandleInitializerDelegateWrapper(
+                IntPtr /* void* */ state,
+                IntPtr /* const char* */ initializerName,
+                IntPtr /* const OrtValue* */ initializerValue,
+                IntPtr /* const OrtExternalInitializerInfo* */ externalInfo,
+                out IntPtr /* OrtExternalInitializerInfo** */ newExternalInfo)
+            {
+                newExternalInfo = IntPtr.Zero;
+
+                try
+                {
+
+                    HandleInitializerConnector connector = (HandleInitializerConnector)GCHandle.FromIntPtr(state).Target;
+                    string utf8InitializerName = NativeOnnxValueHelper.StringFromNativeUtf8(initializerName);
+                    IReadOnlyOrtValue readOnlyInitializerValue = new OrtValue(initializerValue, owned: false);
+                    IReadOnlyExternalInitializerInfo readOnlyExternalInfo = null;
+
+                    if (externalInfo != IntPtr.Zero)
+                    {
+                        readOnlyExternalInfo = new OrtExternalInitializerInfo(externalInfo, ownsHandle: false);
+                    }
+
+                    // Call user's delegate, which may return the new location of the initializer.
+                    OrtExternalInitializerInfo optionalNewExternalInfo = connector._csharpDelegate(
+                        utf8InitializerName, readOnlyInitializerValue, readOnlyExternalInfo);
+
+                    if (optionalNewExternalInfo != null)
+                    {
+                        // Delegate returned info about a new location for the initializer.
+                        // Can't guarantee that the new external info returned by user's delegate is not referenced
+                        // by other C# code. ORT expects to own the new external info, so create a copy here and
+                        // give it to ORT.
+                        string newFilePath = optionalNewExternalInfo.GetFilePath();
+                        byte[] newFilePathBytes = NativeOnnxValueHelper.GetPlatformSerializedString(newFilePath);
+
+                        IntPtr status = NativeMethods.OrtCreateExternalInitializerInfo(
+                            newFilePathBytes,
+                            optionalNewExternalInfo.GetFileOffset(),
+                            (UIntPtr)optionalNewExternalInfo.GetByteSize(),
+                            out newExternalInfo);
+
+                        if (status != IntPtr.Zero)
+                        {
+                            return status;
+                        }
+                    }
+                    else
+                    {
+                        // User's delegate did not return a new location for the initializer. ORT will store initializer
+                        // within the generated compiled model.
+                        newExternalInfo = IntPtr.Zero;
+                    }
+                }
+                catch (Exception ex)
+                {
+                    var error = $"The C# HandleInitializer delegate threw an exception: {ex.Message}";
+                    IntPtr status = NativeMethods.OrtCreateStatus((uint)ErrorCode.Fail,
+                                            NativeOnnxValueHelper.StringToZeroTerminatedUtf8(error));
+                    return status;
+                }
+
+                return IntPtr.Zero;
+            }
+        }
+
+        /// <summary>
+        /// Disposable class that stores resources for a delegate provided by the user.
+        /// </summary>
+        /// <typeparam name="Connector">The type of the connector class (e.g., WriteBufferConnector)</typeparam>
+        /// <typeparam name="Delegate">The type of the native delegate.</typeparam>
+        private class DelegateResources<Connector, Delegate> : IDisposable
+            where Connector : class
+            where Delegate : class
+        {
+            public DelegateResources(Connector connector, Delegate @delegate)
+            {
+                _connector = connector;
+                _delegate = @delegate;
+                _connectorHandle = GCHandle.Alloc(_connector);
+                _delegateHandle = GCHandle.Alloc(_delegate);
+            }
+
+            public IntPtr GetFunctionPointerForDelegate()
+            {
+                return Marshal.GetFunctionPointerForDelegate(_delegate);
+            }
+
+            public IntPtr GetConnectorHandleAsPointer()
+            {
+                return GCHandle.ToIntPtr(_connectorHandle);
+            }
+
+            #region IDispose implementation
+            public void Dispose()
+            {
+                DisposeImpl();
+                GC.SuppressFinalize(this);
+            }
+
+            protected virtual void DisposeImpl()
+            {
+                if (_disposed)
+                {
+                    return;
+                }
+
+                if (_connectorHandle.IsAllocated)
+                {
+                    _connectorHandle.Free();
+                    _connector = null;
+                }
+
+                if (_delegateHandle.IsAllocated)
+                {
+                    _delegateHandle.Free();
+                    _delegate = null;
+                }
+
+                _disposed = true;
+            }
+
+            ~DelegateResources()
+            {
+                DisposeImpl();
+            }
+            #endregion
+
+            private Connector _connector = null;
+            private Delegate _delegate = null;
+            private GCHandle _connectorHandle = default;
+            private GCHandle _delegateHandle = default;
+            private bool _disposed = false;
+        }
+        #endregion
+
         internal IntPtr Handle => handle;
 
 
@@ -146,7 +427,22 @@ protected override bool ReleaseHandle()
         {
             NativeMethods.CompileApi.OrtReleaseModelCompilationOptions(handle);
             handle = IntPtr.Zero;
+
+            _writeBufferDelegateState?.Dispose();
+            _handleInitializerDelegateState?.Dispose();
             return true;
         }
+
+        /// <summary>
+        /// Stores delegate state for the "write buffer" delegate.
+        /// </summary>
+        private DelegateResources<WriteBufferConnector, NativeMethods.DOrtWriteBufferDelegate>
+            _writeBufferDelegateState = null;
+
+        /// <summary>
+        /// Stores delegate state for the "handle initializer" delegate.
+        /// </summary>
+        private DelegateResources<HandleInitializerConnector, NativeMethods.DOrtHandleInitializerDataDelegate>
+            _handleInitializerDelegateState = null;
     }
 }

csharp/src/Microsoft.ML.OnnxRuntime/NativeCompileApiMethods.shared.cs

@@ -21,6 +21,9 @@ public struct OrtCompileApi
         public IntPtr ModelCompilationOptions_SetEpContextEmbedMode;
         public IntPtr CompileModel;
         public IntPtr ModelCompilationOptions_SetFlags;
+        public IntPtr ModelCompilationOptions_SetEpContextBinaryInformation;
+        public IntPtr ModelCompilationOptions_SetOutputModelWriteFunc;
+        public IntPtr ModelCompilationOptions_SetOutputModelHandleInitializerFunc;
     }
 
     internal class NativeMethods
@@ -101,6 +104,31 @@ public DOrtModelCompilationOptions_SetOutputModelExternalInitializersFile
             uint flags);
         public DOrtModelCompilationOptions_SetFlags OrtModelCompilationOptions_SetFlags;
 
+        [UnmanagedFunctionPointer(CallingConvention.Winapi)]
+        public delegate IntPtr /* OrtStatus* */ DOrtModelCompilationOptions_SetEpContextBinaryInformation(
+            IntPtr /* OrtModelCompilationOptions* */ options,
+            byte[] /* const ORTCHAR_T* */ outputDirectory,
+            byte[] /* const ORTCHAR_T* */ modelName);
+        public DOrtModelCompilationOptions_SetEpContextBinaryInformation
+                        OrtModelCompilationOptions_SetEpContextBinaryInformation;
+
+        [UnmanagedFunctionPointer(CallingConvention.Winapi)]
+        public delegate IntPtr /* OrtStatus* */ DOrtModelCompilationOptions_SetOutputModelWriteFunc(
+            IntPtr /* OrtModelCompilationOptions* */ options,
+            IntPtr /* DOrtWriteBufferDelegate */ writeFunc,
+            IntPtr /* void* */ state);
+        public DOrtModelCompilationOptions_SetOutputModelWriteFunc
+                        OrtModelCompilationOptions_SetOutputModelWriteFunc;
+
+        [UnmanagedFunctionPointer(CallingConvention.Winapi)]
+        public delegate IntPtr /* OrtStatus* */ DOrtModelCompilationOptions_SetOutputModelHandleInitializerFunc(
+            IntPtr /* OrtModelCompilationOptions* */ options,
+            IntPtr /* DOrtHandleInitializerDataDelegate */ handleInitializerFunc,
+            IntPtr /* void* */ state);
+        public DOrtModelCompilationOptions_SetOutputModelHandleInitializerFunc
+                        OrtModelCompilationOptions_SetOutputModelHandleInitializerFunc;
+
+
         internal NativeMethods(OnnxRuntime.NativeMethods.DOrtGetCompileApi getCompileApi)
         {
 
@@ -161,6 +189,21 @@ internal NativeMethods(OnnxRuntime.NativeMethods.DOrtGetCompileApi getCompileApi
                     _compileApi.ModelCompilationOptions_SetFlags,
                     typeof(DOrtModelCompilationOptions_SetFlags));
 
+            OrtModelCompilationOptions_SetEpContextBinaryInformation =
+                (DOrtModelCompilationOptions_SetEpContextBinaryInformation)Marshal.GetDelegateForFunctionPointer(
+                    _compileApi.ModelCompilationOptions_SetEpContextBinaryInformation,
+                    typeof(DOrtModelCompilationOptions_SetEpContextBinaryInformation));
+
+            OrtModelCompilationOptions_SetOutputModelWriteFunc =
+                (DOrtModelCompilationOptions_SetOutputModelWriteFunc)Marshal.GetDelegateForFunctionPointer(
+                    _compileApi.ModelCompilationOptions_SetOutputModelWriteFunc,
+                    typeof(DOrtModelCompilationOptions_SetOutputModelWriteFunc));
+
+            OrtModelCompilationOptions_SetOutputModelHandleInitializerFunc =
+                (DOrtModelCompilationOptions_SetOutputModelHandleInitializerFunc)Marshal.GetDelegateForFunctionPointer(
+                    _compileApi.ModelCompilationOptions_SetOutputModelHandleInitializerFunc,
+                    typeof(DOrtModelCompilationOptions_SetOutputModelHandleInitializerFunc));
+
         }
     }
 }