saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Expose IOBinding SynchronizeInputs/Outputs via C/C++/C# And Python APIs (#9823)

Browse source 
Add C/C++ APIs for SynchronizeBoundInputs/Outputs
 Add python bindings
 Expose SynchronizeBoundInputs/Outputs to C# API
This commit is contained in:
Dmitri Smirnov 2021-11-22 09:45:31 -08:00 committed by GitHub
parent 8dcd388ec5
commit 567749b2dc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 212 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

97
csharp/src/Microsoft.ML.OnnxRuntime/NativeMethods.shared.cs
View file

@ -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;
/// <summary>
@ -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
/// <summary>
/// Create OrtIoBinding instance that is used to bind memory that is allocated
@ -1081,6 +1096,20 @@ namespace Microsoft.ML.OnnxRuntime
public static DOrtBindInput OrtBindInput;
/// <summary>
/// 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.
/// </summary>
/// <param name="io_binding">instance of OrtIoBinding</param>
/// <returns>An instance of OrtStatus or null</returns>
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /* OrtStatus*/ DOrtSynchronizeBoundInputs(IntPtr /*(OrtIoBinding)*/ io_binding);
public static DOrtSynchronizeBoundInputs OrtSynchronizeBoundInputs;
/// <summary>
/// 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;
/// <summary>
/// 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.
/// </summary>
/// <param name="io_binding">instance of OrtIoBinding</param>
/// <returns>An instance of OrtStatus or null</returns>
[UnmanagedFunctionPointer(CallingConvention.Winapi)]
public delegate IntPtr /* OrtStatus*/ DOrtSynchronizeBoundOutputs(IntPtr /*(OrtIoBinding)*/ io_binding);
public static DOrtSynchronizeBoundOutputs OrtSynchronizeBoundOutputs;
/// <summary>
/// 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
/// <summary>
/// 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
/// <summary>
/// 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)
{

20
csharp/src/Microsoft.ML.OnnxRuntime/OrtIoBinding.shared.cs
View file

@ -8,7 +8,7 @@ using System.Text;
namespace Microsoft.ML.OnnxRuntime
{
/// <summary>
/// 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);
}
/// <summary>
/// Blocks until device completes all preceding requested tasks.
/// Useful for memory synchronization.
/// </summary>
public void SynchronizeBoundInputs()
{
NativeMethods.OrtSynchronizeBoundInputs(handle);
}
/// <summary>
/// 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));
}
/// <summary>
/// Blocks until device completes all preceding requested tasks.
/// Useful for memory synchronization.
/// </summary>
public void SynchronizeBoundOutputs()
{
NativeMethods.OrtSynchronizeBoundOutputs(handle);
}
/// <summary>
/// Internal helper
/// </summary>

5
csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/InferenceTest.cs
View file

@ -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);

4
csharp/test/Microsoft.ML.OnnxRuntime.Tests.Common/OrtIoBindingAllocationTest.cs
View file

@ -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);

20
include/onnxruntime/core/session/onnxruntime_c_api.h
View file

@ -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);
};
/*

2
include/onnxruntime/core/session/onnxruntime_cxx_api.h
View file

@ -871,6 +871,8 @@ struct IoBinding : public Base<OrtIoBinding> {
std::vector<Value> GetOutputValues(Allocator&) const;
void ClearBoundInputs();
void ClearBoundOutputs();
void SynchronizeInputs();
void SynchronizeOutputs();
private:
std::vector<std::string> GetOutputNamesHelper(OrtAllocator*) const;

8
include/onnxruntime/core/session/onnxruntime_cxx_inline.h
View file

@ -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_));
}

22
onnxruntime/core/session/onnxruntime_c_api.cc
View file

@ -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<const ::OrtValue*>(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)

3
onnxruntime/core/session/ort_apis.h
View file

@ -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

6
onnxruntime/python/onnxruntime_inference_collection.py
View file

@ -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.

12
onnxruntime/python/onnxruntime_pybind_iobinding.cc
View file

@ -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<int64_t>& 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();
})

41
onnxruntime/test/python/onnxruntime_test_python_iobinding.py
View file

@ -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()))

6
onnxruntime/test/shared_lib/test_inference.cc
View file

@ -1053,15 +1053,17 @@ TEST(CApiTest, io_binding_cuda) {
Ort::Value bound_y = Ort::Value::CreateTensor(info_cuda, reinterpret_cast<float*>(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<float, 3 * 2> y_values_0;
cudaMemcpy(y_values_0.data(), output_data.get(), sizeof(float) * y_values_0.size(), cudaMemcpyDeviceToHost);