Expose knobs to create and share (CPU) allocators across sessions in C# and Python

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

@@ -182,6 +182,11 @@ public struct OrtApi
         public IntPtr SetGlobalInterOpNumThreads;
         public IntPtr SetGlobalSpinControl;
         public IntPtr AddInitializer;
+        public IntPtr CreateEnvWithCustomLoggerAndGlobalThreadPools;
+        public IntPtr OrtSessionOptionsAppendExecutionProvider_CUDA;
+        public IntPtr SetGlobalDenormalAsZero;
+        public IntPtr CreateArenaCfg;
+        public IntPtr ReleaseArenaCfg;
     }
 
     internal static class NativeMethods
@@ -260,6 +265,9 @@ static NativeMethods()
             OrtRunOptionsSetTerminate = (DOrtRunOptionsSetTerminate)Marshal.GetDelegateForFunctionPointer(api_.RunOptionsSetTerminate, typeof(DOrtRunOptionsSetTerminate));
             OrtRunOptionsUnsetTerminate = (DOrtRunOptionsUnsetTerminate)Marshal.GetDelegateForFunctionPointer(api_.RunOptionsUnsetTerminate, typeof(DOrtRunOptionsUnsetTerminate));
 
+            OrtCreateArenaCfg = (DOrtCreateArenaCfg)Marshal.GetDelegateForFunctionPointer(api_.CreateArenaCfg, typeof(DOrtCreateArenaCfg));
+            OrtReleaseArenaCfg = (DOrtReleaseArenaCfg)Marshal.GetDelegateForFunctionPointer(api_.ReleaseArenaCfg, typeof(DOrtReleaseArenaCfg));
+            OrtReleaseAllocator = (DOrtReleaseAllocator)Marshal.GetDelegateForFunctionPointer(api_.ReleaseAllocator, typeof(DOrtReleaseAllocator));
             OrtCreateMemoryInfo = (DOrtCreateMemoryInfo)Marshal.GetDelegateForFunctionPointer(api_.CreateMemoryInfo, typeof(DOrtCreateMemoryInfo));
             OrtCreateCpuMemoryInfo = (DOrtCreateCpuMemoryInfo)Marshal.GetDelegateForFunctionPointer(api_.CreateCpuMemoryInfo, typeof(DOrtCreateCpuMemoryInfo));
             OrtReleaseMemoryInfo = (DOrtReleaseMemoryInfo)Marshal.GetDelegateForFunctionPointer(api_.ReleaseMemoryInfo, typeof(DOrtReleaseMemoryInfo));
@@ -310,7 +318,6 @@ static NativeMethods()
             OrtGetTensorShapeElementCount = (DOrtGetTensorShapeElementCount)Marshal.GetDelegateForFunctionPointer(api_.GetTensorShapeElementCount, typeof(DOrtGetTensorShapeElementCount));
             OrtReleaseValue = (DOrtReleaseValue)Marshal.GetDelegateForFunctionPointer(api_.ReleaseValue, typeof(DOrtReleaseValue));
 
-
             OrtSessionGetModelMetadata = (DOrtSessionGetModelMetadata)Marshal.GetDelegateForFunctionPointer(api_.SessionGetModelMetadata, typeof(DOrtSessionGetModelMetadata));
             OrtModelMetadataGetProducerName = (DOrtModelMetadataGetProducerName)Marshal.GetDelegateForFunctionPointer(api_.ModelMetadataGetProducerName, typeof(DOrtModelMetadataGetProducerName));
             OrtModelMetadataGetGraphName = (DOrtModelMetadataGetGraphName)Marshal.GetDelegateForFunctionPointer(api_.ModelMetadataGetGraphName, typeof(DOrtModelMetadataGetGraphName));
@@ -664,6 +671,25 @@ IntPtr[] outputValues /* An array of output value pointers. Array must be alloca
         public delegate IntPtr /*(OrtStatus*)*/DOrtAllocatorGetInfo(IntPtr /*(const OrtAllocator*)*/ ptr, out IntPtr /*(const struct OrtMemoryInfo**)*/info);
         public static DOrtAllocatorGetInfo OrtAllocatorGetInfo;
 
+        /// <summary>
+        /// Create an instance of arena configuration which will be used to create an arena based allocator
+        /// </summary>
+        /// <param name="maxMemory">TODO</param>
+        /// <param name="arenaExtendStrategy">TODO</param>
+        /// <param name="initialChunkSizeBytes">TODO</param>
+        /// <param name="maxDeadBytesPerChunk">TODO</param>
+        public delegate IntPtr /*(OrtStatus*)*/ DOrtCreateArenaCfg(UIntPtr /*(size_t)*/ maxMemory, int /*(int)*/ arenaExtendStrategy,
+                                                                  int /*(int)*/ initialChunkSizeBytes, int /*(int)*/ maxDeadBytesPerChunk,
+                                                                  out IntPtr /*(OrtArenaCfg**)*/ arenaCfg);
+        public static DOrtCreateArenaCfg OrtCreateArenaCfg;
+
+        /// <summary>
+        /// Destroy an instance of an arena configuration instance
+        /// </summary>
+        /// <param name="arenaCfg">arena configuration instance to be destroyed</param>
+        public delegate void DOrtReleaseArenaCfg(IntPtr /*(OrtArenaCfg*)*/ arenaCfg);
+        public static DOrtReleaseArenaCfg OrtReleaseArenaCfg;
+
         /// <summary>
         /// Create an instance of allocator according to mem_info
         /// </summary>
@@ -817,13 +843,15 @@ IntPtr[] outputValues /* An array of output value pointers. Array must be alloca
 
         /// <summary>
         /// Creates an allocator instance and registers it with the env to enable
-        ///sharing between multiple sessions that use the same env instance.
-        ///Lifetime of the created allocator will be valid for the duration of the environment.
-        ///Returns an error if an allocator with the same OrtMemoryInfo is already registered.
-        /// </summary>
-        /// <param name="mem_info">must be non-null</param>
-        /// <param name="arena_cfg">if nullptr defaults will be used</param>
-        public delegate void DOrtCreateAndRegisterAllocator(IntPtr /*(OrtIoBinding)*/ io_binding);
+        /// sharing between multiple sessions that use the same env instance.
+        /// Lifetime of the created allocator will be valid for the duration of the environment.
+        /// Returns an error if an allocator with the same OrtMemoryInfo is already registered.
+        /// <param name="env">Native OrtEnv instance</param>
+        /// <param name="memInfo">Native OrtMemoryInfo instance</param>
+        /// <param name="arenaCfg">Native OrtArenaCfg instance</param>
+        public delegate IntPtr /*(OrtStatus*)*/ DOrtCreateAndRegisterAllocator(IntPtr /*(OrtEnv*)*/ env,
+                                                                               IntPtr /*(const OrtMemoryInfo*)*/ memInfo,
+                                                                               IntPtr/*(const OrtArenaCfg*)*/ arenaCfg);
         public static DOrtCreateAndRegisterAllocator OrtCreateAndRegisterAllocator;
 
         /// <summary>

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

@@ -30,7 +30,7 @@ public enum LogLevel
     public enum OrtLanguageProjection
     {
         ORT_PROJECTION_C = 0,
-        ORT_PROJECTION_CPLUSPLUS = 1 ,
+        ORT_PROJECTION_CPLUSPLUS = 1,
         ORT_PROJECTION_CSHARP = 2,
         ORT_PROJECTION_PYTHON = 3,
         ORT_PROJECTION_JAVA = 4,
@@ -101,6 +101,15 @@ public void DisableTelemetryEvents()
             NativeApiStatus.VerifySuccess(NativeMethods.OrtDisableTelemetryEvents(Handle));
         }
 
+        /// <summary>
+        /// Create and register an arena based allocator to the OrtEnv instance
+        /// so as to enable sharing across all sessions using the OrtEnv instance
+        /// </summary>
+        public void CreateAndRegisterAllocator(OrtMemoryInfo memInfo, OrtArenaCfg arenaCfg)
+        {
+            NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateAndRegisterAllocator(Handle, memInfo.Pointer, arenaCfg.Pointer));
+        }
+
         #endregion
 
         #region SafeHandle

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

@@ -27,6 +27,50 @@ public enum OrtMemType
         Default = 0,                        // the default allocator for execution provider
     }
 
+    /// <summary>
+    /// This class encapsulates an arena configuration information that will be used to define the behavior
+    /// of an arena based allocator
+    /// </summary>
+    public class OrtArenaCfg : SafeHandle
+    {
+        /// <summary>
+        /// Create an instance of arena configuration which will be used to create an arena based allocator
+        /// </summary>
+        /// <param name="maxMemory">TODO</param>
+        /// <param name="arenaExtendStrategy">TODO</param>
+        /// <param name="initialChunkSizeBytes">TODO</param>
+        /// <param name="maxDeadBytesPerChunk">TODO</param>
+        public OrtArenaCfg(uint maxMemory, int arenaExtendStrategy, int initialChunkSizeBytes, int maxDeadBytesPerChunk)
+            : base(IntPtr.Zero, true)
+        {
+            NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateArenaCfg((UIntPtr)maxMemory,
+                                                                           arenaExtendStrategy,
+                                                                           initialChunkSizeBytes,
+                                                                           maxDeadBytesPerChunk,
+                                                                           out handle));
+        }
+
+        public override bool IsInvalid { get { return handle == IntPtr.Zero; } }
+
+        internal IntPtr Pointer
+        {
+            get
+            {
+                return handle;
+            }
+        }
+
+        #region SafeHandle
+        protected override bool ReleaseHandle()
+        {
+            NativeMethods.OrtReleaseArenaCfg(handle);
+            handle = IntPtr.Zero;
+            return true;
+        }
+        #endregion
+
+    }
+
     /// <summary>
     /// This class encapsulates and most of the time owns the underlying native OrtMemoryInfo instance.
     /// Instance returned from OrtAllocator will not own OrtMemoryInfo, the class must be disposed

csharp/test/Microsoft.ML.OnnxRuntime.Tests/InferenceTest.cs

@@ -1892,9 +1892,9 @@ private void TestWeightSharingBetweenSessions()
             var dims = new long[] { 3, 2 };
             var dataBuffer = new float[] { 1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F };
             var dataHandle = GCHandle.Alloc(dataBuffer, GCHandleType.Pinned);
-            
+
             try
-            {            
+            {
                 unsafe
                 {
                     float* p = (float*)dataHandle.AddrOfPinnedObject();
@@ -1927,9 +1927,6 @@ private void TestWeightSharingBetweenSessions()
                         container.Add(NamedOnnxValue.CreateFromTensor<float>(name, tensor));
                     }
 
-                    ReadOnlySpan<int> expectedOutputDimensions = new int[] { 1, 1000, 1, 1 };
-                    string[] expectedOutputNames = new string[] { "Y" };
-
                     // Run inference with named inputs and outputs created with in Run()
                     using (var results = session.Run(container))  // results is an IReadOnlyList<NamedOnnxValue> container
                     {
@@ -1956,6 +1953,83 @@ private void TestWeightSharingBetweenSessions()
             }
         }
 
+        [Fact]
+        private void TestSharedAllocatorUsingCreateAndRegisterAllocator()
+        {
+            string modelPath = Path.Combine(Directory.GetCurrentDirectory(), "mul_1.onnx");
+
+            using (var memInfo = new OrtMemoryInfo(OrtMemoryInfo.allocatorCPU,
+                                                   OrtAllocatorType.ArenaAllocator, 0, OrtMemType.Default))
+            using (var arenaCfg = new OrtArenaCfg(0, -1, -1, -1))
+            {
+                var env = OrtEnv.Instance();
+                // Create and register the arena based allocator
+                env.CreateAndRegisterAllocator(memInfo, arenaCfg);
+
+                using (var sessionOptions = new SessionOptions())
+                {
+                    // Key must match kOrtSessionOptionsConfigUseEnvAllocators in onnxruntime_session_options_config_keys.h
+                    sessionOptions.AddSessionConfigEntry("session.use_env_allocators", "1");
+
+                    // Create two sessions to share the allocator
+                    // Create a thrid session that DOES NOT use the allocator in the environment
+                    using (var session1 = new InferenceSession(modelPath, sessionOptions))
+                    using (var session2 = new InferenceSession(modelPath, sessionOptions))
+                    using (var session3 = new InferenceSession(modelPath)) // Use the default SessionOptions instance
+                    {
+                        // Input data
+                        var inputDims = new long[] { 3, 2 };
+                        var input = new float[] { 1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F };
+
+                        // Output data
+                        int[] outputDims = { 3, 2 };
+                        float[] output = { 1.0F, 4.0F, 9.0F, 16.0F, 25.0F, 36.0F };
+
+                        // Run inference on all three models
+                        var inputMeta = session1.InputMetadata;
+                        var container = new List<NamedOnnxValue>();
+
+                        foreach (var name in inputMeta.Keys)
+                        {
+                            Assert.Equal(typeof(float), inputMeta[name].ElementType);
+                            Assert.True(inputMeta[name].IsTensor);
+                            var tensor = new DenseTensor<float>(input, inputMeta[name].Dimensions);
+                            container.Add(NamedOnnxValue.CreateFromTensor<float>(name, tensor));
+                        }
+
+                        // Run inference with named inputs and outputs created with in Run()
+                        using (var results = session1.Run(container))  // results is an IReadOnlyList<NamedOnnxValue> container
+                        {
+                            foreach (var r in results)
+                            {
+                                validateRunResultData(r.AsTensor<float>(), output, outputDims);
+                            }
+                        }
+
+                        // Run inference with named inputs and outputs created with in Run()
+                        using (var results = session2.Run(container))  // results is an IReadOnlyList<NamedOnnxValue> container
+                        {
+                            foreach (var r in results)
+                            {
+                                validateRunResultData(r.AsTensor<float>(), output, outputDims);
+                            }
+                        }
+
+                        // Run inference with named inputs and outputs created with in Run()
+                        using (var results = session3.Run(container))  // results is an IReadOnlyList<NamedOnnxValue> container
+                        {
+                            foreach (var r in results)
+                            {
+                                validateRunResultData(r.AsTensor<float>(), output, outputDims);
+                            }
+                        }
+                    }
+                }
+
+
+            }
+        }
+
         [DllImport("kernel32", SetLastError = true)]
         static extern IntPtr LoadLibrary(string lpFileName);
 
@@ -2034,7 +2108,6 @@ static float[] LoadTensorFromFile(string filename, bool skipheader = true)
             return tensorData.ToArray();
         }
 
-
         private enum TensorElementType
         {
             Float = 1,

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -1103,6 +1103,20 @@ struct OrtApi {
    * and that's recommended because turning this option on may hurt model accuracy.
    */
   ORT_API2_STATUS(SetGlobalDenormalAsZero, _Inout_ OrtThreadingOptions* tp_options);
+
+  /**
+  * Use this API to create the configuration of an arena that can eventually be used to define 
+  * an arena based allocator's behavior
+  * max_mem : use 0 to allow ORT to choose the default
+  * arena_extend_strategy :  use -1 to allow ORT to choose the default, 0 = kNextPowerOfTwo, 1 = kSameAsRequested
+  * initial_chunk_size_bytes : use -1 to allow ORT to choose the default
+  * max_dead_bytes_per_chunk : use -1 to allow ORT to choose the default
+  * See ONNX_Runtime_Perf_Tuning.md for details on what these mean and how to choose these values
+  */
+  ORT_API2_STATUS(CreateArenaCfg, _In_ size_t max_mem, int arena_extend_strategy, int initial_chunk_size_bytes,
+                  int max_dead_bytes_per_chunk, _Outptr_ OrtArenaCfg** out);
+
+  ORT_CLASS_RELEASE(ArenaCfg);
 };
 
 /*

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -98,6 +98,7 @@ ORT_DEFINE_RELEASE(Value);
 ORT_DEFINE_RELEASE(ModelMetadata);
 ORT_DEFINE_RELEASE(ThreadingOptions);
 ORT_DEFINE_RELEASE(IoBinding);
+ORT_DEFINE_RELEASE(ArenaCfg);
 
 // This is used internally by the C++ API. This is the common base class used by the wrapper objects.
 template <typename T>
@@ -252,7 +253,6 @@ struct SessionOptions : Base<OrtSessionOptions> {
   SessionOptions& AddConfigEntry(const char* config_key, const char* config_value);
   SessionOptions& AddInitializer(const char* name, const OrtValue* ort_val);
   OrtStatus* OrtSessionOptionsAppendExecutionProvider_CUDA(OrtSessionOptions* options, OrtCUDAProviderOptions* cuda_options);
-
 };
 
 struct ModelMetadata : Base<OrtModelMetadata> {
@@ -479,6 +479,11 @@ struct IoBinding : public Base<OrtIoBinding> {
   void ClearBoundOutputs();
 };
 
+struct ArenaCfg : Base<OrtArenaCfg> {
+  explicit ArenaCfg(std::nullptr_t) {}
+  ArenaCfg(size_t max_mem, int arena_extend_strategy, int initial_chunk_size_bytes, int max_dead_bytes_per_chunk);
+};
+
 //
 // Custom OPs (only needed to implement custom OPs)
 //

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -286,6 +286,10 @@ inline void IoBinding::ClearBoundOutputs() {
   GetApi().ClearBoundOutputs(p_);
 }
 
+inline ArenaCfg::ArenaCfg(size_t max_mem, int arena_extend_strategy, int initial_chunk_size_bytes, int max_dead_bytes_per_chunk) {
+  ThrowOnError(GetApi().CreateArenaCfg(max_mem, arena_extend_strategy, initial_chunk_size_bytes, max_dead_bytes_per_chunk, &p_));
+}
+
 inline Env::Env(OrtLoggingLevel logging_level, _In_ const char* logid) {
   ThrowOnError(GetApi().CreateEnv(logging_level, logid, &p_));
   if (strcmp(logid, "onnxruntime-node") == 0) {
@@ -477,7 +481,7 @@ inline SessionOptions& SessionOptions::AddInitializer(const char* name, const Or
   return *this;
 }
 
-inline OrtStatus* SessionOptions::OrtSessionOptionsAppendExecutionProvider_CUDA(OrtSessionOptions * options, OrtCUDAProviderOptions * cuda_options) {
+inline OrtStatus* SessionOptions::OrtSessionOptionsAppendExecutionProvider_CUDA(OrtSessionOptions* options, OrtCUDAProviderOptions* cuda_options) {
   ThrowOnError(GetApi().OrtSessionOptionsAppendExecutionProvider_CUDA(options, cuda_options));
   return nullptr;
 }

onnxruntime/__init__.py

@@ -12,7 +12,8 @@
 
 from onnxruntime.capi._pybind_state import get_all_providers, get_available_providers, get_device, set_seed, \
     RunOptions, SessionOptions, set_default_logger_severity, enable_telemetry_events, disable_telemetry_events, \
-    NodeArg, ModelMetadata, GraphOptimizationLevel, ExecutionMode, ExecutionOrder, OrtDevice, SessionIOBinding
+    NodeArg, ModelMetadata, GraphOptimizationLevel, ExecutionMode, ExecutionOrder, OrtDevice, SessionIOBinding, \
+    OrtAllocatorType, OrtMemType
 
 try:
     from onnxruntime.capi._pybind_state import set_cuda_mem_limit, set_cuda_device_id

onnxruntime/core/session/inference_session.cc

@@ -1098,6 +1098,7 @@ common::Status InferenceSession::Initialize() {
     std::string use_env_allocators = session_options_.GetConfigOrDefault(kOrtSessionOptionsConfigUseEnvAllocators,
                                                                          "0");
     if (use_env_allocators == "1") {
+      LOGS(*session_logger_, INFO) << "This session will use the allocator registered with the environment.";
       UpdateProvidersWithSharedAllocators();
     }