How to create Tensor on GPU directly

[messense]

Currently Tensor::from_array allocates the value on CPU, is there an equivalent method to allocate the value on GPU device directly?

[decahedron1]

Tensor::new

[messense]

Given this example

let allocator = Allocator::new(
	&session,
	MemoryInfo::new(AllocationDevice::CUDAPinned, 0, AllocatorType::Device, MemoryType::CPUInput)?
)?;

let mut img_input = Tensor::<f32>::new(&allocator, [1, 128, 128, 3])?;

How to set the actual data for img_input?

[decahedron1]

ONNX Runtime doesn't provide APIs for that. You can grab the pointer to the data with Tensor::data_ptr_mut, cast it to a more concrete data type (in the case of CUDA, this would be a CUDA buffer) and write the data that way.

[messense]

That's unfortunate, can it be done by modify Tensor::from_array and pass in a memory_info instead?

ort/src/value/impl_tensor/create.rs

Line 180 in ce5aaba

let memory_info = MemoryInfo::new_cpu(AllocatorType::Arena, MemoryType::Default)?;

[decahedron1]

The memory_info specifies where the data buffer itself resides, it won't copy data from CPU to GPU.

Maybe I/O binding is a better fit?

[messense]

Yeah I'm actually try to use I/O binding, but couldn't find an example that actually binds input/output.
It's a bit unclear how to use the I/O binding API correctly to reduce copies between CPU & GPU.