Skip to content

Commit 2d36f04

Browse files
adrianlizarragatianleiwu
authored andcommitted
Compile API: output model and initializer stream write functions (#25455)
### Description - Adds `ModelCompilationOptions_SetOutputModelWriteFunc` to the compile API to allow writing the output model ONNX bytes to a user-provided write function (i.e., for streaming). - Adds `ModelCompilationOptions_SetOutputModelHandleInitializerFunc` to the compile API to allow the user to write individual initializers to some destination. Also allows specifying if an initializer should be embedded within the ONNX model or written to a custom file. - Adds C++, Python, and C# bindings for the new APIs. A follow-up PR adds a write function for EPContext node binary data: #25471 ### Example `ModelCompilationOptions_SetOutputModelWriteFunc`: https://github.com/microsoft/onnxruntime/blob/c62ed23c328cbbfefd3083c1f7a6ced604772c19/onnxruntime/test/providers/qnn/qnn_ep_context_test.cc#L2075-L2131 `ModelCompilationOptions_SetOutputModelHandleInitializerFunc`: https://github.com/microsoft/onnxruntime/blob/c62ed23c328cbbfefd3083c1f7a6ced604772c19/onnxruntime/test/providers/qnn/qnn_ep_context_test.cc#L2160-L2292 ### Motivation and Context Add output streaming capabilities when saving compiled models.
1 parent ab71f1e commit 2d36f04

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

42 files changed

+2694
-348
lines changed

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

Lines changed: 349 additions & 23 deletions
Large diffs are not rendered by default.

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ public struct OrtCompileApi
2323
public IntPtr ModelCompilationOptions_SetFlags;
2424
public IntPtr ModelCompilationOptions_SetEpContextBinaryInformation;
2525
public IntPtr ModelCompilationOptions_SetGraphOptimizationLevel;
26+
public IntPtr ModelCompilationOptions_SetOutputModelWriteFunc;
27+
public IntPtr ModelCompilationOptions_SetOutputModelGetInitializerLocationFunc;
2628
}
2729

2830
internal class NativeMethods
@@ -118,6 +120,22 @@ public DOrtModelCompilationOptions_SetEpContextBinaryInformation
118120
public DOrtModelCompilationOptions_SetGraphOptimizationLevel
119121
OrtModelCompilationOptions_SetGraphOptimizationLevel;
120122

123+
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
124+
public delegate IntPtr /* OrtStatus* */ DOrtModelCompilationOptions_SetOutputModelWriteFunc(
125+
IntPtr /* OrtModelCompilationOptions* */ options,
126+
IntPtr /* DOrtWriteBufferDelegate */ writeFunc,
127+
IntPtr /* void* */ state);
128+
public DOrtModelCompilationOptions_SetOutputModelWriteFunc
129+
OrtModelCompilationOptions_SetOutputModelWriteFunc;
130+
131+
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
132+
public delegate IntPtr /* OrtStatus* */ DOrtModelCompilationOptions_SetOutputModelGetInitializerLocationFunc(
133+
IntPtr /* OrtModelCompilationOptions* */ options,
134+
IntPtr /* DOrtHandleInitializerDataDelegate */ handleInitializerFunc,
135+
IntPtr /* void* */ state);
136+
public DOrtModelCompilationOptions_SetOutputModelGetInitializerLocationFunc
137+
OrtModelCompilationOptions_SetOutputModelGetInitializerLocationFunc;
138+
121139
internal NativeMethods(OnnxRuntime.NativeMethods.DOrtGetCompileApi getCompileApi)
122140
{
123141

@@ -188,6 +206,17 @@ internal NativeMethods(OnnxRuntime.NativeMethods.DOrtGetCompileApi getCompileApi
188206
_compileApi.ModelCompilationOptions_SetGraphOptimizationLevel,
189207
typeof(DOrtModelCompilationOptions_SetGraphOptimizationLevel));
190208

209+
OrtModelCompilationOptions_SetOutputModelWriteFunc =
210+
(DOrtModelCompilationOptions_SetOutputModelWriteFunc)Marshal.GetDelegateForFunctionPointer(
211+
_compileApi.ModelCompilationOptions_SetOutputModelWriteFunc,
212+
typeof(DOrtModelCompilationOptions_SetOutputModelWriteFunc));
213+
214+
OrtModelCompilationOptions_SetOutputModelGetInitializerLocationFunc =
215+
(DOrtModelCompilationOptions_SetOutputModelGetInitializerLocationFunc)Marshal.
216+
GetDelegateForFunctionPointer(
217+
_compileApi.ModelCompilationOptions_SetOutputModelGetInitializerLocationFunc,
218+
typeof(DOrtModelCompilationOptions_SetOutputModelGetInitializerLocationFunc));
219+
191220
}
192221
}
193222
}

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

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -450,6 +450,7 @@ public struct OrtApi
450450

451451
public IntPtr Graph_GetModelMetadata;
452452
public IntPtr GetModelCompatibilityForEpDevices;
453+
public IntPtr CreateExternalInitializerInfo;
453454
}
454455

455456
internal static class NativeMethods
@@ -787,9 +788,35 @@ static NativeMethods()
787788
api_.SessionOptionsSetEpSelectionPolicyDelegate,
788789
typeof(DSessionOptionsSetEpSelectionPolicyDelegate));
789790

791+
OrtReleaseExternalInitializerInfo =
792+
(DOrtReleaseExternalInitializerInfo)Marshal.GetDelegateForFunctionPointer(
793+
api_.ReleaseExternalInitializerInfo,
794+
typeof(DOrtReleaseExternalInitializerInfo));
795+
796+
OrtExternalInitializerInfo_GetFilePath =
797+
(DOrtExternalInitializerInfo_GetFilePath)Marshal.GetDelegateForFunctionPointer(
798+
api_.ExternalInitializerInfo_GetFilePath,
799+
typeof(DOrtExternalInitializerInfo_GetFilePath));
800+
801+
OrtExternalInitializerInfo_GetFileOffset =
802+
(DOrtExternalInitializerInfo_GetFileOffset)Marshal.GetDelegateForFunctionPointer(
803+
api_.ExternalInitializerInfo_GetFileOffset,
804+
typeof(DOrtExternalInitializerInfo_GetFileOffset));
805+
806+
OrtExternalInitializerInfo_GetByteSize =
807+
(DOrtExternalInitializerInfo_GetByteSize)Marshal.GetDelegateForFunctionPointer(
808+
api_.ExternalInitializerInfo_GetByteSize,
809+
typeof(DOrtExternalInitializerInfo_GetByteSize));
810+
790811
OrtGetModelCompatibilityForEpDevices = (DOrtGetModelCompatibilityForEpDevices)Marshal.GetDelegateForFunctionPointer(
791812
api_.GetModelCompatibilityForEpDevices,
792813
typeof(DOrtGetModelCompatibilityForEpDevices));
814+
815+
OrtCreateExternalInitializerInfo =
816+
(DOrtCreateExternalInitializerInfo)Marshal.GetDelegateForFunctionPointer(
817+
api_.CreateExternalInitializerInfo,
818+
typeof(DOrtCreateExternalInitializerInfo));
819+
793820
}
794821

795822
internal class NativeLib
@@ -2382,6 +2409,70 @@ out IntPtr lora_adapter
23822409
public delegate ref CompileApi.OrtCompileApi DOrtGetCompileApi();
23832410
#endif
23842411
public static DOrtGetCompileApi OrtGetCompileApi;
2412+
2413+
/// <summary>
2414+
/// Delegate called by ORT to write a buffer (ONNX model bytes) to a custom destination (e.g., file or stream).
2415+
/// </summary>
2416+
/// <param name="state">State that was provided in when the delegate was registered.</param>
2417+
/// <param name="buffer">The buffer to write.</param>
2418+
/// <param name="bufferNumBytes">The size of the buffer in bytes.</param>
2419+
/// <returns>OrtStatus*</returns>
2420+
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
2421+
public delegate IntPtr /* OrtStatus* */ DOrtWriteBufferToDestinationDelegate(
2422+
IntPtr /* void* */ state,
2423+
IntPtr /* const void* */ buffer,
2424+
UIntPtr /* size_t */ bufferNumBytes
2425+
);
2426+
2427+
/// <summary>
2428+
/// Function called by ORT to allow user to specify how an initializer should be saved while compiling
2429+
/// a model, that is, either written to an external file or stored within the model. ORT calls this function
2430+
/// for every initializer.
2431+
/// </summary>
2432+
/// <param name="state">State that was provided when the delegate was registered.</param>
2433+
/// <param name="initializerName">The initializer's name.</param>
2434+
/// <param name="initializerValue">The OrtValue containing the initializer's data, type, and shape</param>
2435+
/// <param name="externalInfo">The original initializer's location in an external file, or NULL.</param>
2436+
/// <param name="newExternalInfo">Output parameter set to a new OrtExternalInitializerInfo instance
2437+
/// indicating the location where the function implementation stored the initializer data. If the function
2438+
/// implementation sets `newExternalInfo` to NULL, ORT stores the initializer within the generated model.</param>
2439+
/// <returns></returns>
2440+
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
2441+
public delegate IntPtr /* OrtStatus* */ DOrtGetInitializerLocationDelegate(
2442+
IntPtr /* void* */ state,
2443+
IntPtr /* const char* */ initializerName,
2444+
IntPtr /* const OrtValue* */ initializerValue,
2445+
IntPtr /* const OrtExternalInitializerInfo* */ externalInfo,
2446+
out IntPtr /* OrtExternalInitializerInfo** */ newExternalInfo
2447+
);
2448+
2449+
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
2450+
public delegate void DOrtReleaseExternalInitializerInfo(IntPtr /* OrtExternalInitializerInfo* */ info);
2451+
2452+
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
2453+
public delegate IntPtr /* OrtStatus* */ DOrtCreateExternalInitializerInfo(
2454+
byte[] /* const ORTCHAR_T* */ filePath,
2455+
long /* int64_t */ fileOffset,
2456+
UIntPtr /* size_t */ byteSize,
2457+
out IntPtr /* OrtExternalInitializerInfo** */ outInfo);
2458+
2459+
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
2460+
public delegate IntPtr /* const ORTCHAR_T* */ DOrtExternalInitializerInfo_GetFilePath(
2461+
IntPtr /* const OrtExternalInitializerInfo* */ info);
2462+
2463+
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
2464+
public delegate long /* int64_t */ DOrtExternalInitializerInfo_GetFileOffset(
2465+
IntPtr /* const OrtExternalInitializerInfo* */ info);
2466+
2467+
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
2468+
public delegate UIntPtr /* size_t */ DOrtExternalInitializerInfo_GetByteSize(
2469+
IntPtr /* const OrtExternalInitializerInfo* */ info);
2470+
2471+
public static DOrtReleaseExternalInitializerInfo OrtReleaseExternalInitializerInfo;
2472+
public static DOrtCreateExternalInitializerInfo OrtCreateExternalInitializerInfo;
2473+
public static DOrtExternalInitializerInfo_GetFilePath OrtExternalInitializerInfo_GetFilePath;
2474+
public static DOrtExternalInitializerInfo_GetFileOffset OrtExternalInitializerInfo_GetFileOffset;
2475+
public static DOrtExternalInitializerInfo_GetByteSize OrtExternalInitializerInfo_GetByteSize;
23852476
#endregion
23862477

23872478
#region Auto EP API related

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,45 @@ internal static byte[] GetPlatformSerializedString(string str)
150150
else
151151
return StringToZeroTerminatedUtf8(str);
152152
}
153+
154+
/// <summary>
155+
/// Converts a null-terminated path string that is pointed to by the given IntPtr handle into
156+
/// a C# UTF-16 string.
157+
/// </summary>
158+
/// <remarks>A path string on Windows is utf-16, but utf-8 on other operating systems.</remarks>
159+
/// <param name="strPtr"></param>
160+
/// <returns></returns>
161+
internal static string StringFromNativePathString(IntPtr strPtr)
162+
{
163+
if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
164+
{
165+
if (strPtr == IntPtr.Zero)
166+
{
167+
return string.Empty;
168+
}
169+
170+
// Get length of utf16 string by checking for two 0 bytes in a row.
171+
int length = 0;
172+
while (Marshal.ReadInt16(strPtr, length * 2) != 0)
173+
{
174+
length += 1;
175+
}
176+
177+
if (length == 0)
178+
{
179+
return string.Empty;
180+
}
181+
182+
unsafe
183+
{
184+
return System.Text.Encoding.Unicode.GetString((byte*)strPtr, length * 2);
185+
}
186+
}
187+
else
188+
{
189+
return StringFromNativeUtf8(strPtr);
190+
}
191+
}
153192
}
154193

155194
// Guards an array of disposable objects on stack and disposes them in reverse order
Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
5+
namespace Microsoft.ML.OnnxRuntime
6+
{
7+
using System;
8+
using System.Diagnostics;
9+
using System.Runtime.InteropServices;
10+
11+
/// <summary>
12+
/// Class to that stores information about the file location where an "external" initializer is stored.
13+
/// </summary>
14+
/// <see cref="OrtModelCompilationOptions.HandleInitializerDelegate"/>
15+
public class OrtExternalInitializerInfo : SafeHandle, IReadOnlyExternalInitializerInfo
16+
{
17+
// Set to false when constructed with an externally managed constant handle owned by ORT.
18+
private readonly bool _ownsHandle = true;
19+
20+
/// <summary>
21+
/// Create a new OrtExternalInitializerInfo instance.
22+
/// </summary>
23+
/// <param name="filePath">The path to the file that stores the initializer data.</param>
24+
/// <param name="fileOffset">The byte offset in the file where the data is stored.</param>
25+
/// <param name="byteSize">The size of the data (in bytes) within the file.</param>
26+
public OrtExternalInitializerInfo(string filePath, long fileOffset, long byteSize)
27+
: base(IntPtr.Zero, ownsHandle: true)
28+
{
29+
var platformFilePath = NativeOnnxValueHelper.GetPlatformSerializedString(filePath);
30+
NativeApiStatus.VerifySuccess(
31+
NativeMethods.OrtCreateExternalInitializerInfo(platformFilePath, fileOffset, (UIntPtr)byteSize, out handle));
32+
_ownsHandle = true;
33+
}
34+
35+
/// <summary>
36+
/// Create a new OrtExternalInitializerInfo instance from an existing native OrtExternalInitializerInfo handle.
37+
/// </summary>
38+
/// <param name="constHandle">Native OrtExternalInitializerInfo handle.</param>
39+
/// <param name="ownsHandle">True if the OrtExternalInitializerInfo instance owns the native handle.
40+
/// Defaults to false.</param>
41+
internal OrtExternalInitializerInfo(IntPtr constHandle, bool ownsHandle = false)
42+
: base(IntPtr.Zero, ownsHandle)
43+
{
44+
Debug.Assert(constHandle != IntPtr.Zero);
45+
SetHandle(constHandle);
46+
_ownsHandle = ownsHandle;
47+
}
48+
49+
/// <summary>
50+
/// Get the file path to the file that store's the initializer's data.
51+
/// </summary>
52+
/// <remarks>
53+
/// The path is relative to the filesystem directory where the ONNX model was stored.
54+
/// </remarks>
55+
/// <returns>The file path.</returns>
56+
public string GetFilePath()
57+
{
58+
IntPtr filePathPtr = NativeMethods.OrtExternalInitializerInfo_GetFilePath(handle);
59+
if (filePathPtr == IntPtr.Zero)
60+
{
61+
return string.Empty;
62+
}
63+
64+
return NativeOnnxValueHelper.StringFromNativePathString(filePathPtr);
65+
}
66+
67+
/// <summary>
68+
/// Get the byte offset within the file where the initializer's data is stored.
69+
/// </summary>
70+
/// <returns>The file offset location.</returns>
71+
public long GetFileOffset()
72+
{
73+
return NativeMethods.OrtExternalInitializerInfo_GetFileOffset(handle);
74+
}
75+
76+
/// <summary>
77+
/// Get the size in bytes of the initializer's data within the file.
78+
/// </summary>
79+
/// <returns>The size in bytes of the initializer data.</returns>
80+
public long GetByteSize()
81+
{
82+
UIntPtr byteSize = NativeMethods.OrtExternalInitializerInfo_GetByteSize(handle);
83+
return checked((long)byteSize);
84+
}
85+
86+
/// <summary>
87+
/// Indicates whether the native handle is invalid.
88+
/// </summary>
89+
public override bool IsInvalid { get { return handle == IntPtr.Zero; } }
90+
91+
/// <summary>
92+
/// Release the native instance of OrtExternalInitializerInfo if we own it.
93+
/// </summary>
94+
/// <returns>true on success and false on error.</returns>
95+
protected override bool ReleaseHandle()
96+
{
97+
if (!_ownsHandle)
98+
{
99+
// Return false to indicate an error.
100+
// ReleaseHandle() should not be called on a const handle that this class does not own.
101+
return false;
102+
}
103+
104+
NativeMethods.OrtReleaseExternalInitializerInfo(handle);
105+
handle = IntPtr.Zero;
106+
return true;
107+
}
108+
}
109+
110+
/// <summary>
111+
/// Interface for all readonly methods implemented by OrtExternalInitializerInfo.
112+
/// </summary>
113+
public interface IReadOnlyExternalInitializerInfo
114+
{
115+
/// <summary>
116+
/// Get the file path to the file that store's the initializer's data.
117+
/// </summary>
118+
/// <remarks>
119+
/// The path is relative to the filesystem directory where the ONNX model was stored.
120+
/// </remarks>
121+
/// <returns>The file path.</returns>
122+
string GetFilePath();
123+
124+
/// <summary>
125+
/// Get the byte offset within the file where the initializer's data is stored.
126+
/// </summary>
127+
/// <returns>The file offset location.</returns>
128+
long GetFileOffset();
129+
130+
/// <summary>
131+
/// Get the size in bytes of the initializer's data within the file.
132+
/// </summary>
133+
/// <returns>The size in bytes of the initializer data.</returns>
134+
long GetByteSize();
135+
}
136+
}

0 commit comments

Comments
 (0)