Model chaining with dynamic output shapes

[cNoNim]

Hey everyone!

I've been working on chaining multiple ONNX models in a pipeline (like Model A → Model B → Model C), and ran into a limitation with IOBinding when dealing with dynamic output shapes. I wanted to discuss the problem and a potential solution I've prototyped.

Problem

When you want to pass outputs from one model as inputs to another, ONNX Runtime already supports this well... if you know the output shapes in advance:

auto shared_tensor = Value::CreateTensor(memory_info, data_ptr, size, known_shape, shape_len);

IoBinding binding1(session1);
IoBinding binding2(session2);

binding1.BindOutput("output", shared_tensor);
binding2.BindInput("input", shared_tensor);

session1.Run(Ort::RunOptions(), binding1);
session2.Run(Ort::RunOptions(), binding2);

Just works!

But what if the output shape is dynamic? Then you have to use BindOutputToDevice, which only specifies the device, not the actual tensor. And now you're doing this:

session1.Run(Ort::RunOptions(), binding1);
std::vector<Ort::Value> outputs = binding1.GetOutputValues();
binding2.BindInput("input", outputs[0]);
session2.Run(Ort::RunOptions(), binding2);

This works, but the problem is:

• Multiple API calls across FFI boundary (GetOutputValues → BindInput) for every inference in language bindings
• Can't set up the pipeline once

Solution

What if we could set up the pipeline once even with dynamic shapes, and have the "copy output to input" happen automatically?

IoBinding binding1(session1);
IoBinding binding2(session2);

MemoryInfo memory_info = ...;
binding1.BindOutput("output", memory_info);

std::shared_ptr<OrtValue> handle;
binding1.GetOutputHandle("output", handle);

binding2.BindInput("input", handle);

session1.Run(Ort::RunOptions(), binding1);  // Output gets allocated dynamically
session2.Run(Ort::RunOptions(), binding2);  // Input automatically uses updated data

Implementation

I've made a small prototype that adds this functionality. Here's the basic idea:

When you get an output handle:

1. We create a std::shared_ptr<OrtValue>
2. Store it in a separate array (shared_outputs_) in IOBinding

After Run() completes we call SyncSharedOutputs() internally which automatically copies outputs into the shared handles

When you bind a shared handle as input:

1. We store the std::shared_ptr<OrtValue> in another array (shared_inputs_)
2. Keep track of which input index it corresponds to

Before Run() executes we call SyncSharedInputs() internally which automatically copies from shared handles into the input bindings

Questions

1. Does this seem like the right approach? Better alternatives? Maybe there's a simpler way I'm missing?
2. Shared ownership: The prototype uses std::shared_ptr<OrtValue> - is this acceptable? In C API I provide OrtOutputHandle wrapper type.
3. API design: Is it okay to have GetOutputHandle() as a new method and overloaded BindInput() that accepts shared_ptr?

Notes

• After GetOutputHandle() is called, subsequent BindOutput() calls on that output maintain the created handle in correct state (sync it automatically)
• When you bind a shared handle as input, subsequent BindInput() calls on the same input name also update the handle. This might need a different strategy or an unbinding method?
• If at Run() time a bound shared handle is not allocated yet, it's an error