add more to iobinding

yuzawa-san · Jan 9, 2024 · 083b5d0 · 083b5d0
1 parent 58f3edb
commit 083b5d0
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 2 deletions.
diff --git a/src/main/java/com/jyuzawa/onnxruntime/ApiImpl.java b/src/main/java/com/jyuzawa/onnxruntime/ApiImpl.java
@@ -147,6 +147,8 @@ final class ApiImpl implements Api {
     final SessionOptionsAppendExecutionProvider_OpenVINO SessionOptionsAppendExecutionProvider_OpenVINO;
     final SessionOptionsAppendExecutionProvider_ROCM SessionOptionsAppendExecutionProvider_ROCM;
     final SessionOptionsAppendExecutionProvider_TensorRT_V2 SessionOptionsAppendExecutionProvider_TensorRT_V2;
+    final SynchronizeBoundInputs SynchronizeBoundInputs;
+    final SynchronizeBoundOutputs SynchronizeBoundOutputs;
     final UpdateCUDAProviderOptions UpdateCUDAProviderOptions;
     final UpdateDnnlProviderOptions UpdateDnnlProviderOptions;
     final UpdateTensorRTProviderOptions UpdateTensorRTProviderOptions;
@@ -290,6 +292,8 @@ final class ApiImpl implements Api {
                 OrtApi.SessionOptionsAppendExecutionProvider_ROCM(memorySegment, memorySession);
         this.SessionOptionsAppendExecutionProvider_TensorRT_V2 =
                 OrtApi.SessionOptionsAppendExecutionProvider_TensorRT_V2(memorySegment, memorySession);
+        this.SynchronizeBoundInputs = OrtApi.SynchronizeBoundInputs(memorySegment, memorySession);
+        this.SynchronizeBoundOutputs = OrtApi.SynchronizeBoundOutputs(memorySegment, memorySession);
         this.UpdateCUDAProviderOptions = OrtApi.UpdateCUDAProviderOptions(memorySegment, memorySession);
         this.UpdateDnnlProviderOptions = OrtApi.UpdateDnnlProviderOptions(memorySegment, memorySession);
         this.UpdateTensorRTProviderOptions = OrtApi.UpdateTensorRTProviderOptions(memorySegment, memorySession);

diff --git a/src/main/java/com/jyuzawa/onnxruntime/IoBinding.java b/src/main/java/com/jyuzawa/onnxruntime/IoBinding.java
@@ -36,6 +36,10 @@ public interface IoBinding extends AutoCloseable {
      */
     IoBinding setRunTag(String runTag);
 
+    IoBinding synchronizeBoundInputs();
+
+    IoBinding synchronizeBoundOutputs();
+
     NamedCollection<OnnxValue> getInputs();
 
     NamedCollection<OnnxValue> getOutputs();

diff --git a/src/main/java/com/jyuzawa/onnxruntime/IoBindingImpl.java b/src/main/java/com/jyuzawa/onnxruntime/IoBindingImpl.java
@@ -95,8 +95,8 @@ static final class Builder implements IoBinding.Builder {
         public Builder(SessionImpl session) {
             this.api = session.api;
             this.session = session;
-            this.inputs = new ArrayList<>();
-            this.outputs = new ArrayList<>();
+            this.inputs = new ArrayList<>(session.inputs.size());
+            this.outputs = new ArrayList<>(session.outputs.size());
         }
 
         @Override
@@ -178,4 +178,16 @@ public NamedCollection<OnnxValue> getInputs() {
     public NamedCollection<OnnxValue> getOutputs() {
         return outputs;
     }
+
+    @Override
+    public IoBinding synchronizeBoundInputs() {
+        api.checkStatus(api.SynchronizeBoundInputs.apply(ioBinding));
+        return this;
+    }
+
+    @Override
+    public IoBinding synchronizeBoundOutputs() {
+        api.checkStatus(api.SynchronizeBoundOutputs.apply(ioBinding));
+        return this;
+    }
 }
diff --git a/src/test/java/com/jyuzawa/onnxruntime/SessionTest.java b/src/test/java/com/jyuzawa/onnxruntime/SessionTest.java
@@ -984,7 +984,9 @@ public void ioBindingTest() throws IOException {
                     ThreadLocalRandom.current().nextInt()
                 };
                 inputBuf.clear().put(rawInput);
+                txn.synchronizeBoundInputs();
                 txn.run();
+                txn.synchronizeBoundOutputs();
                 outputBuf.rewind().get(rawOutput);
                 assertArrayEquals(rawInput, rawOutput);
             }