diff --git a/csharp/src/Microsoft.ML.OnnxRuntime/NativeMethods.shared.cs b/csharp/src/Microsoft.ML.OnnxRuntime/NativeMethods.shared.cs index e708950d24..09b3cc48c9 100644 --- a/csharp/src/Microsoft.ML.OnnxRuntime/NativeMethods.shared.cs +++ b/csharp/src/Microsoft.ML.OnnxRuntime/NativeMethods.shared.cs @@ -223,6 +223,18 @@ namespace Microsoft.ML.OnnxRuntime public IntPtr GetSparseTensorValues; public IntPtr GetSparseTensorIndicesTypeShape; public IntPtr GetSparseTensorIndices; + public IntPtr HasValue; + public IntPtr KernelContext_GetGPUComputeStream; + public IntPtr GetTensorMemoryInfo; + public IntPtr GetExecutionProviderApi; + public IntPtr SessionOptionsSetCustomCreateThreadFn; + public IntPtr SessionOptionsSetCustomThreadCreationOptions; + public IntPtr SessionOptionsSetCustomJoinThreadFn; + public IntPtr SetGlobalCustomCreateThreadFn; + public IntPtr SetGlobalCustomThreadCreationOptions; + public IntPtr SetGlobalCustomJoinThreadFn; + public IntPtr SynchronizeBoundInputs; + public IntPtr SynchronizeBoundOutputs; } internal static class NativeMethods @@ -328,12 +340,15 @@ namespace Microsoft.ML.OnnxRuntime OrtCreateIoBinding = (DOrtCreateIoBinding)Marshal.GetDelegateForFunctionPointer(api_.CreateIoBinding, typeof(DOrtCreateIoBinding)); OrtReleaseIoBinding = (DOrtReleaseIoBinding)Marshal.GetDelegateForFunctionPointer(api_.ReleaseIoBinding, typeof(DOrtReleaseIoBinding)); OrtBindInput = (DOrtBindInput)Marshal.GetDelegateForFunctionPointer(api_.BindInput, typeof(DOrtBindInput)); + OrtSynchronizeBoundInputs = (DOrtSynchronizeBoundInputs)Marshal.GetDelegateForFunctionPointer(api_.SynchronizeBoundInputs, typeof(DOrtSynchronizeBoundInputs)); OrtBindOutput = (DOrtBindOutput)Marshal.GetDelegateForFunctionPointer(api_.BindOutput, typeof(DOrtBindOutput)); OrtBindOutputToDevice = (DOrtBindOutputToDevice)Marshal.GetDelegateForFunctionPointer(api_.BindOutputToDevice, typeof(DOrtBindOutputToDevice)); + OrtSynchronizeBoundOutputs = (DOrtSynchronizeBoundOutputs)Marshal.GetDelegateForFunctionPointer(api_.SynchronizeBoundOutputs, typeof(DOrtSynchronizeBoundOutputs)); OrtGetBoundOutputNames = (DOrtGetBoundOutputNames)Marshal.GetDelegateForFunctionPointer(api_.GetBoundOutputNames, typeof(DOrtGetBoundOutputNames)); OrtGetBoundOutputValues = (DOrtGetBoundOutputValues)Marshal.GetDelegateForFunctionPointer(api_.GetBoundOutputValues, typeof(DOrtGetBoundOutputValues)); OrtClearBoundInputs = (DOrtClearBoundInputs)Marshal.GetDelegateForFunctionPointer(api_.ClearBoundInputs, typeof(DOrtClearBoundInputs)); OrtClearBoundOutputs = (DOrtClearBoundOutputs)Marshal.GetDelegateForFunctionPointer(api_.ClearBoundOutputs, typeof(DOrtClearBoundOutputs)); + OrtTensorAt = (DOrtTensorAt)Marshal.GetDelegateForFunctionPointer(api_.TensorAt, typeof(DOrtTensorAt)); OrtCreateAndRegisterAllocator = (DOrtCreateAndRegisterAllocator)Marshal.GetDelegateForFunctionPointer(api_.CreateAndRegisterAllocator, typeof(DOrtCreateAndRegisterAllocator)); OrtSetLanguageProjection = (DOrtSetLanguageProjection)Marshal.GetDelegateForFunctionPointer(api_.SetLanguageProjection, typeof(DOrtSetLanguageProjection)); @@ -482,7 +497,7 @@ namespace Microsoft.ML.OnnxRuntime byte[] modelPath, IntPtr /* (OrtSessionOptions*) */sessopnOptions, out IntPtr /**/ session); - + public static DOrtCreateSession OrtCreateSession; /// @@ -500,7 +515,7 @@ namespace Microsoft.ML.OnnxRuntime IntPtr /* (OrtSessionOptions*) */sessionOptions, IntPtr /* (OrtPrepackedWeightsContainer*) */prepackedWeightsContainer, out IntPtr /* (OrtSession**) */ session); - + public static DOrtCreateSessionWithPrepackedWeightsContainer OrtCreateSessionWithPrepackedWeightsContainer; [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -542,7 +557,7 @@ namespace Microsoft.ML.OnnxRuntime UIntPtr outputCount, IntPtr[] outputValues /* An array of output value pointers. Array must be allocated by the caller */ ); - + public static DOrtRun OrtRun; [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -551,28 +566,28 @@ namespace Microsoft.ML.OnnxRuntime IntPtr /*(OrtSessionRunOptions*)*/ runOptions, // can not be null IntPtr /*(const OrtIoBinding*)*/ io_binding ); - + public static DOrtRunWithBinding OrtRunWithBinding; [UnmanagedFunctionPointer(CallingConvention.Winapi)] public delegate IntPtr /*(OrtStatus*)*/ DOrtSessionGetInputCount( IntPtr /*(OrtSession*)*/ session, out UIntPtr count); - + public static DOrtSessionGetInputCount OrtSessionGetInputCount; [UnmanagedFunctionPointer(CallingConvention.Winapi)] public delegate IntPtr /*(OrtStatus*)*/ DOrtSessionGetOutputCount( IntPtr /*(OrtSession*)*/ session, out UIntPtr count); - + public static DOrtSessionGetOutputCount OrtSessionGetOutputCount; [UnmanagedFunctionPointer(CallingConvention.Winapi)] public delegate IntPtr /*(OrtStatus*)*/ DOrtSessionGetOverridableInitializerCount( IntPtr /*(OrtSession*)*/ session, out UIntPtr count); - + public static DOrtSessionGetOverridableInitializerCount OrtSessionGetOverridableInitializerCount; [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -581,7 +596,7 @@ namespace Microsoft.ML.OnnxRuntime UIntPtr index, IntPtr /*(OrtAllocator*)*/ allocator, out IntPtr /*(char**)*/name); - + public static DOrtSessionGetInputName OrtSessionGetInputName; [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -590,7 +605,7 @@ namespace Microsoft.ML.OnnxRuntime UIntPtr index, IntPtr /*(OrtAllocator*)*/ allocator, out IntPtr /*(char**)*/name); - + public static DOrtSessionGetOutputName OrtSessionGetOutputName; [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -598,7 +613,7 @@ namespace Microsoft.ML.OnnxRuntime IntPtr /*(const OrtSession*)*/ session, IntPtr /*(OrtAllocator*)*/ allocator, out IntPtr /*(char**)*/profile_file); - + public static DOrtSessionEndProfiling OrtSessionEndProfiling; [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -607,7 +622,7 @@ namespace Microsoft.ML.OnnxRuntime UIntPtr index, IntPtr /*(OrtAllocator*)*/ allocator, out IntPtr /*(char**)*/name); - + public static DOrtSessionGetOverridableInitializerName OrtSessionGetOverridableInitializerName; [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -615,7 +630,7 @@ namespace Microsoft.ML.OnnxRuntime IntPtr /*(const OrtSession*)*/ session, UIntPtr index, out IntPtr /*(struct OrtTypeInfo**)*/ typeInfo); - + public static DOrtSessionGetInputTypeInfo OrtSessionGetInputTypeInfo; [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -623,7 +638,7 @@ namespace Microsoft.ML.OnnxRuntime IntPtr /*(const OrtSession*)*/ session, UIntPtr index, out IntPtr /* (struct OrtTypeInfo**)*/ typeInfo); - + public static DOrtSessionGetOutputTypeInfo OrtSessionGetOutputTypeInfo; [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -631,7 +646,7 @@ namespace Microsoft.ML.OnnxRuntime IntPtr /*(const OrtSession*)*/ session, UIntPtr index, out IntPtr /* (struct OrtTypeInfo**)*/ typeInfo); - + public static DOrtSessionGetOverridableInitializerTypeInfo OrtSessionGetOverridableInitializerTypeInfo; // release the typeinfo using OrtReleaseTypeInfo @@ -854,9 +869,9 @@ namespace Microsoft.ML.OnnxRuntime public static DOrtAddInitializer OrtAddInitializer; -#endregion + #endregion -#region RunOptions API + #region RunOptions API [UnmanagedFunctionPointer(CallingConvention.Winapi)] public delegate IntPtr /*(OrtStatus*)*/ DOrtCreateRunOptions(out IntPtr /* OrtRunOptions** */ runOptions); @@ -900,9 +915,9 @@ namespace Microsoft.ML.OnnxRuntime public delegate IntPtr /*(OrtStatus*)*/ DOrtRunOptionsUnsetTerminate(IntPtr /* OrtRunOptions* */ options); public static DOrtRunOptionsUnsetTerminate OrtRunOptionsUnsetTerminate; -#endregion + #endregion -#region Allocator/MemoryInfo API + #region Allocator/MemoryInfo API [UnmanagedFunctionPointer(CallingConvention.Winapi)] public delegate IntPtr /* (OrtStatus*)*/ DOrtCreateMemoryInfo( @@ -1041,9 +1056,9 @@ namespace Microsoft.ML.OnnxRuntime public static DOrtAllocatorFree OrtAllocatorFree; -#endregion Allocator/MemoryInfo API + #endregion Allocator/MemoryInfo API -#region IoBinding API + #region IoBinding API /// /// Create OrtIoBinding instance that is used to bind memory that is allocated @@ -1081,6 +1096,20 @@ namespace Microsoft.ML.OnnxRuntime public static DOrtBindInput OrtBindInput; + /// + /// The API calls Sync() on all EP providers present. This blocks until the device has completed + /// all preceding requested tasks. This is necessary when memory synchronization is required. + /// For example, the memory bound to an input is likely to be on a different CUDA stream. + /// For some scenarios and devices this may be a no-op, use + /// your best judgment. + /// + /// instance of OrtIoBinding + /// An instance of OrtStatus or null + [UnmanagedFunctionPointer(CallingConvention.Winapi)] + public delegate IntPtr /* OrtStatus*/ DOrtSynchronizeBoundInputs(IntPtr /*(OrtIoBinding)*/ io_binding); + + public static DOrtSynchronizeBoundInputs OrtSynchronizeBoundInputs; + /// /// Bind OrtValue to the model output with the specified name /// If binding with the specified name already exists, it will be replaced @@ -1109,6 +1138,18 @@ namespace Microsoft.ML.OnnxRuntime public static DOrtBindOutputToDevice OrtBindOutputToDevice; + /// + /// The API calls Sync() on all EP providers present. This blocks until the device has completed + /// all preceding requested tasks. This is necessary when memory synchronization is required. + /// For some scenarios and devices this may be a no-op, use your best judgment. + /// + /// instance of OrtIoBinding + /// An instance of OrtStatus or null + [UnmanagedFunctionPointer(CallingConvention.Winapi)] + public delegate IntPtr /* OrtStatus*/ DOrtSynchronizeBoundOutputs(IntPtr /*(OrtIoBinding)*/ io_binding); + + public static DOrtSynchronizeBoundOutputs OrtSynchronizeBoundOutputs; + /// /// The function will return all bound output names in the order they were bound. /// It is the same order that the output values will be returned after RunWithBinding() is used. @@ -1205,9 +1246,9 @@ namespace Microsoft.ML.OnnxRuntime public static DOrtSetLanguageProjection OrtSetLanguageProjection; -#endregion IoBinding API + #endregion IoBinding API -#region ModelMetadata API + #region ModelMetadata API /// /// Gets the ModelMetadata associated with an InferenceSession @@ -1326,9 +1367,9 @@ namespace Microsoft.ML.OnnxRuntime public static DOrtReleaseModelMetadata OrtReleaseModelMetadata; -#endregion ModelMetadata API + #endregion ModelMetadata API -#region Tensor/OnnxValue API + #region Tensor/OnnxValue API [UnmanagedFunctionPointer(CallingConvention.Winapi)] public delegate IntPtr /*(OrtStatus*)*/ DOrtGetValue(IntPtr /*(OrtValue*)*/ value, @@ -1337,7 +1378,7 @@ namespace Microsoft.ML.OnnxRuntime out IntPtr /*(OrtValue**)*/ outputValue); public static DOrtGetValue OrtGetValue; - + [UnmanagedFunctionPointer(CallingConvention.Winapi)] public delegate IntPtr /*(OrtStatus*)*/ DOrtGetValueType(IntPtr /*(OrtValue*)*/ value, out IntPtr /*(OnnxValueType*)*/ onnxtype); @@ -1485,9 +1526,9 @@ namespace Microsoft.ML.OnnxRuntime public static DOrtReleaseValue OrtReleaseValue; -#endregion + #endregion -#region Misc API + #region Misc API /// /// Queries all the execution providers supported in the native onnxruntime shared library @@ -1527,7 +1568,7 @@ namespace Microsoft.ML.OnnxRuntime public static DOrtReleasePrepackedWeightsContainer OrtReleasePrepackedWeightsContainer; -#endregion + #endregion public static byte[] GetPlatformSerializedString(string str) { diff --git a/csharp/src/Microsoft.ML.OnnxRuntime/OrtIoBinding.shared.cs b/csharp/src/Microsoft.ML.OnnxRuntime/OrtIoBinding.shared.cs index 84bda820ee..40549a6848 100644 --- a/csharp/src/Microsoft.ML.OnnxRuntime/OrtIoBinding.shared.cs +++ b/csharp/src/Microsoft.ML.OnnxRuntime/OrtIoBinding.shared.cs @@ -8,7 +8,7 @@ using System.Text; namespace Microsoft.ML.OnnxRuntime { /// - /// This class enable to bind inputs and outputs to pre-allocated + /// This class enables binding of inputs and/or outputs to pre-allocated /// memory. This enables interesting scenarios. For example, if your input /// already resides in some pre-allocated memory like GPU, you can bind /// that piece of memory to an input name and shape and onnxruntime will use that as input. @@ -87,6 +87,15 @@ namespace Microsoft.ML.OnnxRuntime BindInputOrOutput(name, fixedValue.Value.Handle, true); } + /// + /// Blocks until device completes all preceding requested tasks. + /// Useful for memory synchronization. + /// + public void SynchronizeBoundInputs() + { + NativeMethods.OrtSynchronizeBoundInputs(handle); + } + /// /// Bind model output to an OrtValue as Tensor with a given type and shape. An instance of OrtMemoryAllocaiton /// owns the memory and should be alive for the time of execution.The size of the allocation can not be less than required @@ -133,6 +142,15 @@ namespace Microsoft.ML.OnnxRuntime NativeApiStatus.VerifySuccess(NativeMethods.OrtBindOutputToDevice(handle, pinnedName.Pointer, memInfo.Pointer)); } + /// + /// Blocks until device completes all preceding requested tasks. + /// Useful for memory synchronization. + /// + public void SynchronizeBoundOutputs() + { + NativeMethods.OrtSynchronizeBoundOutputs(handle); + } + /// /// Internal helper /// diff --git a/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/InferenceTest.cs b/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/InferenceTest.cs index 892bfb2865..b9ed2d4f26 100644 --- a/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/InferenceTest.cs +++ b/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/InferenceTest.cs @@ -1670,8 +1670,10 @@ namespace Microsoft.ML.OnnxRuntime.Tests { ioBinding.BindInput(inputName, fixeInputBuffer); ioBinding.BindOutput(outputName, fixedOutputBuffer); + ioBinding.SynchronizeBoundInputs(); using (var outputs = session.RunWithBindingAndNames(runOptions, ioBinding)) { + ioBinding.SynchronizeBoundOutputs(); Assert.Equal(1, outputs.Count); var output = outputs.First(); Assert.Equal(outputName, output.Name); @@ -1687,9 +1689,10 @@ namespace Microsoft.ML.OnnxRuntime.Tests { ioBinding.BindInput(inputName, fixedInputBuffer); ioBinding.BindOutputToDevice(outputName, allocator.Info); - + ioBinding.SynchronizeBoundInputs(); using (var outputs = session.RunWithBindingAndNames(runOptions, ioBinding)) { + ioBinding.SynchronizeBoundOutputs(); Assert.Equal(1, outputs.Count); var output = outputs.First(); Assert.Equal(outputName, output.Name); diff --git a/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/OrtIoBindingAllocationTest.cs b/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/OrtIoBindingAllocationTest.cs index 84cea5ea90..7c9fcfe348 100644 --- a/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/OrtIoBindingAllocationTest.cs +++ b/csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/OrtIoBindingAllocationTest.cs @@ -73,8 +73,10 @@ namespace Microsoft.ML.OnnxRuntime.Tests { ioBinding.BindInput(inputName, fixedInputBuffer); ioBinding.BindOutput(outputName, Tensors.TensorElementType.Float, outputShape, ortAllocationOutput); + ioBinding.SynchronizeBoundInputs(); using (var outputs = session.RunWithBindingAndNames(runOptions, ioBinding)) { + ioBinding.SynchronizeBoundOutputs(); Assert.Equal(1, outputs.Count); var output = outputs.ElementAt(0); Assert.Equal(outputName, output.Name); @@ -88,8 +90,10 @@ namespace Microsoft.ML.OnnxRuntime.Tests { ioBinding.BindInput(inputName, Tensors.TensorElementType.Float, inputShape, ortAllocationInput); ioBinding.BindOutput(outputName, Tensors.TensorElementType.Float, outputShape, ortAllocationOutput); + ioBinding.SynchronizeBoundInputs(); using (var outputs = session.RunWithBindingAndNames(runOptions, ioBinding)) { + ioBinding.SynchronizeBoundOutputs(); Assert.Equal(1, outputs.Count); var output = outputs.ElementAt(0); Assert.Equal(outputName, output.Name); diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 8b9c2871ff..2ac9ce5b8c 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -3162,6 +3162,26 @@ struct OrtApi { */ ORT_API2_STATUS(SetGlobalCustomJoinThreadFn, _Inout_ OrtThreadingOptions* tp_options, _In_ OrtCustomJoinThreadFn ort_custom_join_thread_fn); /// @} + + /** \brief Synchronize bound inputs. The call may be necessary for some providers, such as cuda, + * in case the system that allocated bound memory operated on a different stream. However, the + * operation is provider specific and could be a no-op. + * + * \param[inout] binding_ptr + * + * * \snippet{doc} snippets.dox OrtStatus Return Value + */ + ORT_API2_STATUS(SynchronizeBoundInputs, _Inout_ OrtIoBinding* binding_ptr); + + /** \brief Synchronize bound outputs. The call may be necessary for some providers, such as cuda, + * in case the system that allocated bound memory operated on a different stream. However, the + * operation is provider specific and could be a no-op. + * + * \param[inout] binding_ptr + * + * * \snippet{doc} snippets.dox OrtStatus Return Value + */ + ORT_API2_STATUS(SynchronizeBoundOutputs, _Inout_ OrtIoBinding* binding_ptr); }; /* diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index 0c2b74b39b..048421099b 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -871,6 +871,8 @@ struct IoBinding : public Base { std::vector GetOutputValues(Allocator&) const; void ClearBoundInputs(); void ClearBoundOutputs(); + void SynchronizeInputs(); + void SynchronizeOutputs(); private: std::vector GetOutputNamesHelper(OrtAllocator*) const; diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h index 982a516cca..1f31dffca8 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h @@ -284,6 +284,14 @@ inline void IoBinding::ClearBoundOutputs() { GetApi().ClearBoundOutputs(p_); } +inline void IoBinding::SynchronizeInputs() { + ThrowOnError(GetApi().SynchronizeBoundInputs(p_)); +} + +inline void IoBinding::SynchronizeOutputs() { + ThrowOnError(GetApi().SynchronizeBoundOutputs(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_)); } diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index 16f8b1f41a..7894eb602b 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -970,6 +970,26 @@ ORT_API(void, OrtApis::ClearBoundOutputs, _Inout_ OrtIoBinding* binding_ptr) { binding_ptr->binding_->ClearOutputs(); } +ORT_API_STATUS_IMPL(OrtApis::SynchronizeBoundInputs, _Inout_ OrtIoBinding* binding_ptr) { + API_IMPL_BEGIN + auto st = binding_ptr->binding_->SynchronizeInputs(); + if (!st.IsOK()) { + return ToOrtStatus(st); + } + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::SynchronizeBoundOutputs, _Inout_ OrtIoBinding* binding_ptr) { + API_IMPL_BEGIN + auto st = binding_ptr->binding_->SynchronizeOutputs(); + if (!st.IsOK()) { + return ToOrtStatus(st); + } + return nullptr; + API_IMPL_END +} + ORT_API_STATUS_IMPL(OrtApis::IsTensor, _In_ const OrtValue* value, _Out_ int* out) { auto v = reinterpret_cast(value); *out = v->IsTensor() ? 1 : 0; @@ -2485,6 +2505,8 @@ static constexpr OrtApi ort_api_1_to_10 = { &OrtApis::SetGlobalCustomCreateThreadFn, &OrtApis::SetGlobalCustomThreadCreationOptions, &OrtApis::SetGlobalCustomJoinThreadFn, + &OrtApis::SynchronizeBoundInputs, + &OrtApis::SynchronizeBoundOutputs }; // Asserts to do a some checks to ensure older Versions of the OrtApi never change (will detect an addition or deletion but not if they cancel out each other) diff --git a/onnxruntime/core/session/ort_apis.h b/onnxruntime/core/session/ort_apis.h index 5053b8ac5f..2345e00ed3 100644 --- a/onnxruntime/core/session/ort_apis.h +++ b/onnxruntime/core/session/ort_apis.h @@ -325,5 +325,6 @@ ORT_API_STATUS_IMPL(SessionOptionsSetCustomJoinThreadFn, _Inout_ OrtSessionOptio ORT_API_STATUS_IMPL(SetGlobalCustomCreateThreadFn, _Inout_ OrtThreadingOptions* tp_options, _In_ OrtCustomCreateThreadFn ort_custom_create_thread_fn); ORT_API_STATUS_IMPL(SetGlobalCustomThreadCreationOptions, _Inout_ OrtThreadingOptions* tp_options, _In_ void* ort_custom_thread_creation_options); ORT_API_STATUS_IMPL(SetGlobalCustomJoinThreadFn, _Inout_ OrtThreadingOptions* tp_options, _In_ OrtCustomJoinThreadFn ort_custom_join_thread_fn); - +ORT_API_STATUS_IMPL(SynchronizeBoundInputs, _Inout_ OrtIoBinding* binding_ptr); +ORT_API_STATUS_IMPL(SynchronizeBoundOutputs, _Inout_ OrtIoBinding* binding_ptr); } // namespace OrtApis diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index 0e92984e01..b0a86bce64 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -447,6 +447,9 @@ class IOBinding: ''' self._iobinding.bind_ortvalue_input(name, ortvalue._ortvalue) + def synchronize_inputs(self): + self._iobinding.synchronize_inputs() + def bind_output(self, name, device_type='cpu', device_id=0, element_type=None, shape=None, buffer_ptr=None): ''' :param name: output name @@ -482,6 +485,9 @@ class IOBinding: ''' self._iobinding.bind_ortvalue_output(name, ortvalue._ortvalue) + def synchronize_outputs(self): + self._iobinding.synchronize_outputs() + def get_outputs(self): ''' Returns the output OrtValues from the Run() that preceded the call. diff --git a/onnxruntime/python/onnxruntime_pybind_iobinding.cc b/onnxruntime/python/onnxruntime_pybind_iobinding.cc index 852dd2a136..fb88844f4f 100644 --- a/onnxruntime/python/onnxruntime_pybind_iobinding.cc +++ b/onnxruntime/python/onnxruntime_pybind_iobinding.cc @@ -87,6 +87,12 @@ void addIoBindingMethods(pybind11::module& m) { throw std::runtime_error("Error when binding input: " + status.ErrorMessage()); } }) + .def("synchronize_inputs", [](SessionIOBinding* io_binding) -> void { + auto status = io_binding->Get()->SynchronizeInputs(); + if (!status.IsOK()) { + throw std::runtime_error("Error when synchronizing bound inputs: " + status.ErrorMessage()); + } + }) // This binds output to a pre-allocated memory as a Tensor .def("bind_output", [](SessionIOBinding* io_binding, const std::string& name, const OrtDevice& device, py::object& element_type, std::vector& shape, int64_t data_ptr) -> void { ORT_ENFORCE(data_ptr != 0, "Pointer to data memory is not valid"); @@ -140,6 +146,12 @@ void addIoBindingMethods(pybind11::module& m) { throw std::runtime_error("Error when binding output: " + status.ErrorMessage()); } }) + .def("synchronize_outputs", [](SessionIOBinding* io_binding) -> void { + auto status = io_binding->Get()->SynchronizeOutputs(); + if (!status.IsOK()) { + throw std::runtime_error("Error when synchronizing bound outputs: " + status.ErrorMessage()); + } + }) .def("clear_binding_inputs", [](SessionIOBinding* io_binding) -> void { io_binding->Get()->ClearInputs(); }) diff --git a/onnxruntime/test/python/onnxruntime_test_python_iobinding.py b/onnxruntime/test/python/onnxruntime_test_python_iobinding.py index d7ead9fbb0..4d2660c053 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_iobinding.py +++ b/onnxruntime/test/python/onnxruntime_test_python_iobinding.py @@ -39,6 +39,9 @@ class TestIOBinding(unittest.TestCase): # Invoke Run session.run_with_iobinding(io_binding) + # Sync if different CUDA streams + io_binding.synchronize_outputs() + # Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here) ort_output = io_binding.copy_outputs_to_cpu()[0] @@ -54,11 +57,17 @@ class TestIOBinding(unittest.TestCase): # Bind input to CUDA io_binding.bind_input('X', 'cuda', 0, np.float32, [3, 2], input.data_ptr()) + # Sync if different CUDA streams + io_binding.synchronize_inputs() + # Bind output to CPU io_binding.bind_output('Y') # Invoke Run session.run_with_iobinding(io_binding) + + # Sync if different CUDA streams + io_binding.synchronize_outputs() # Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here) ort_output = io_binding.copy_outputs_to_cpu()[0] @@ -79,8 +88,14 @@ class TestIOBinding(unittest.TestCase): output = self.create_uninitialized_ortvalue_input_on_gpu() io_binding.bind_output('Y', 'cuda', 0, np.float32, [3, 2], output.data_ptr()) + # Sync if different CUDA streams + io_binding.synchronize_inputs() + # Invoke Run session.run_with_iobinding(io_binding) + + # Sync if different CUDA streams + io_binding.synchronize_outputs() # Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here) ort_output_vals = io_binding.copy_outputs_to_cpu()[0] @@ -104,9 +119,15 @@ class TestIOBinding(unittest.TestCase): # Bind output to CUDA io_binding.bind_output('Y', 'cuda') + # Sync if different CUDA streams + io_binding.synchronize_inputs() + # Invoke Run session.run_with_iobinding(io_binding) + # Sync if different CUDA streams + io_binding.synchronize_outputs() + # This call returns an OrtValue which has data allocated by ORT on CUDA ort_outputs = io_binding.get_outputs() self.assertEqual(len(ort_outputs), 1) @@ -124,10 +145,16 @@ class TestIOBinding(unittest.TestCase): # Change the bound input and validate the results in the same bound OrtValue # Bind alternate input to CUDA io_binding.bind_input('X', 'cuda', 0, np.float32, [3, 2], self.create_ortvalue_alternate_input_on_gpu().data_ptr()) - + + # Sync if different CUDA streams + io_binding.synchronize_inputs() + # Invoke Run session.run_with_iobinding(io_binding) + # Sync if different CUDA streams + io_binding.synchronize_outputs() + # This call returns an OrtValue which has data allocated by ORT on CUDA ort_outputs = io_binding.get_outputs() self.assertEqual(len(ort_outputs), 1) @@ -147,9 +174,15 @@ class TestIOBinding(unittest.TestCase): output_ortvalue = self.create_uninitialized_ortvalue_input_on_gpu() io_binding.bind_ortvalue_output('Y', output_ortvalue) + # Sync if different CUDA streams + io_binding.synchronize_inputs() + # Invoke Run session.run_with_iobinding(io_binding) + # Sync if different CUDA streams + io_binding.synchronize_outputs() + # Inspect contents of output_ortvalue and make sure that it has the right contents self.assertTrue(np.array_equal(self.create_expected_output(), output_ortvalue.numpy())) @@ -157,9 +190,15 @@ class TestIOBinding(unittest.TestCase): input_ortvalue_2 = self.create_ortvalue_alternate_input_on_gpu() io_binding.bind_ortvalue_input('X', input_ortvalue_2) + # Sync if different CUDA streams + io_binding.synchronize_inputs() + # Invoke Run session.run_with_iobinding(io_binding) + # Sync if different CUDA streams + io_binding.synchronize_outputs() + # Inspect contents of output_ortvalue and make sure that it has the right contents self.assertTrue(np.array_equal(self.create_expected_output_alternate(), output_ortvalue.numpy())) diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index 3aca885ea2..0a48df4f0b 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -1053,15 +1053,17 @@ TEST(CApiTest, io_binding_cuda) { Ort::Value bound_y = Ort::Value::CreateTensor(info_cuda, reinterpret_cast(output_data.get()), expected_y.size(), expected_y_shape.data(), expected_y_shape.size()); - // Sychronize to make sure the copy on default stream is done since TensorRT isn't using default stream. - cudaStreamSynchronize(nullptr); Ort::IoBinding binding(session); binding.BindInput("X", bound_x); binding.BindOutput("Y", bound_y); + // Sychronize to make sure the copy on default stream is done since TensorRT isn't using default stream. + binding.SynchronizeInputs(); session.Run(Ort::RunOptions(), binding); + binding.SynchronizeOutputs(); + // Check the values against the bound raw memory (needs copying from device to host first) std::array y_values_0; cudaMemcpy(y_values_0.data(), output_data.get(), sizeof(float) * y_values_0.size(), cudaMemcpyDeviceToHost);