Merge pull request #67 from yuzawa-san/perfier

cpu and memory optimization
yuzawa-san · Nov 19, 2022 · a066012 · a066012
2 parents 40e14ef + d81b8a5
commit a066012
Show file tree

Hide file tree

Showing 11 changed files with 99 additions and 88 deletions.
diff --git a/src/main/java/com/jyuzawa/onnxruntime/ApiImpl.java b/src/main/java/com/jyuzawa/onnxruntime/ApiImpl.java
@@ -102,6 +102,7 @@
 import java.lang.foreign.MemoryAddress;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.MemorySession;
+import java.lang.foreign.SegmentAllocator;
 import java.util.function.Function;
 
 final class ApiImpl implements Api {
@@ -306,19 +307,19 @@ void checkStatus(Addressable rawAddress) {
         throw new OnnxRuntimeException(code, message);
     }
 
-    MemoryAddress create(MemorySession allocator, Function<MemoryAddress, Addressable> constructor) {
+    MemoryAddress create(SegmentAllocator allocator, Function<MemoryAddress, Addressable> constructor) {
         MemorySegment pointer = allocator.allocate(C_POINTER);
         checkStatus(constructor.apply(pointer.address()));
         return pointer.getAtIndex(C_POINTER, 0);
     }
 
-    int extractInt(MemorySession allocator, Function<MemoryAddress, Addressable> method) {
+    int extractInt(SegmentAllocator allocator, Function<MemoryAddress, Addressable> method) {
         MemorySegment pointer = allocator.allocate(C_INT);
         checkStatus(method.apply(pointer.address()));
         return pointer.getAtIndex(C_INT, 0);
     }
 
-    long extractLong(MemorySession allocator, Function<MemoryAddress, Addressable> method) {
+    long extractLong(SegmentAllocator allocator, Function<MemoryAddress, Addressable> method) {
         MemorySegment pointer = allocator.allocate(C_LONG);
         checkStatus(method.apply(pointer.address()));
         return pointer.getAtIndex(C_LONG, 0);

diff --git a/src/main/java/com/jyuzawa/onnxruntime/OnnxMapImpl.java b/src/main/java/com/jyuzawa/onnxruntime/OnnxMapImpl.java
@@ -9,6 +9,7 @@
 import java.lang.foreign.MemoryAddress;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.MemorySession;
+import java.lang.foreign.SegmentAllocator;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.LinkedHashMap;
@@ -49,11 +50,11 @@ static final OnnxValueImpl fromTypeInfo(MapInfoImpl mapInfo) {
         }
     }
 
-    private final T newKeyVector(int size, MemorySession scope) {
+    private final T newKeyVector(int size, SegmentAllocator scope) {
         return keyVectorFactory.apply(TensorInfoImpl.of(mapInfo.getKeyType(), size, scope));
     }
 
-    private final OnnxTensorImpl newValueVector(int size, MemorySession scope) {
+    private final OnnxTensorImpl newValueVector(int size, SegmentAllocator scope) {
         return OnnxTensorImpl.fromTypeInfo(
                 TensorInfoImpl.of(mapInfo.getValueType().getTensorInfo().getType(), size, scope));
     }
@@ -93,7 +94,7 @@ public OnnxTypedMap<String> asStringMap() {
 
     @Override
     public MemoryAddress toNative(
-            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, MemorySession allocator) {
+            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, SegmentAllocator allocator) {
         int size = data.size();
         T keyVector = newKeyVector(size, allocator);
         implodeKeyVector(keyVector, data.keySet());
@@ -105,29 +106,26 @@ public MemoryAddress toNative(
         MemorySegment kvArray = allocator.allocateArray(C_POINTER, 2);
         kvArray.setAtIndex(C_POINTER, 0, keyAddress);
         kvArray.setAtIndex(C_POINTER, 1, valueAddress);
-        MemoryAddress value = api.create(
-                allocator, out -> api.CreateValue.apply(kvArray.address(), 2, OnnxType.MAP.getNumber(), out));
-        allocator.addCloseAction(() -> {
-            api.ReleaseValue.apply(value);
-        });
-        return value;
+        return api.create(allocator, out -> api.CreateValue.apply(kvArray.address(), 2, OnnxType.MAP.getNumber(), out));
     }
 
     @Override
-    public void fromNative(ApiImpl api, MemoryAddress ortAllocator, MemoryAddress address, MemorySession allocator) {
+    public void fromNative(
+            ApiImpl api,
+            MemoryAddress ortAllocator,
+            MemoryAddress address,
+            SegmentAllocator allocator,
+            MemorySession session) {
         MemoryAddress keyAddress = api.create(allocator, out -> api.GetValue.apply(address, 0, ortAllocator, out));
         MemoryAddress valueAddress = api.create(allocator, out -> api.GetValue.apply(address, 1, ortAllocator, out));
         MemoryAddress keyInfo = api.create(allocator, out -> api.GetTensorTypeAndShape.apply(keyAddress, out));
         int size =
                 Math.toIntExact(api.extractLong(allocator, out -> api.GetTensorShapeElementCount.apply(keyInfo, out)));
         T keyVector = newKeyVector(size, allocator);
         OnnxTensorImpl valueVector = newValueVector(size, allocator);
-        keyVector.fromNative(api, ortAllocator, keyAddress, allocator);
-        valueVector.fromNative(api, ortAllocator, valueAddress, allocator);
+        keyVector.fromNative(api, ortAllocator, keyAddress, allocator, session);
+        valueVector.fromNative(api, ortAllocator, valueAddress, allocator, session);
         valueVector.getScalars(explodeKeyVector(keyVector).map(this::set));
-        allocator.addCloseAction(() -> {
-            api.ReleaseValue.apply(address);
-        });
     }
 
     @Override

diff --git a/src/main/java/com/jyuzawa/onnxruntime/OnnxOpaqueImpl.java b/src/main/java/com/jyuzawa/onnxruntime/OnnxOpaqueImpl.java
@@ -6,6 +6,7 @@
 
 import java.lang.foreign.MemoryAddress;
 import java.lang.foreign.MemorySession;
+import java.lang.foreign.SegmentAllocator;
 import java.nio.ByteBuffer;
 
 final class OnnxOpaqueImpl extends OnnxValueImpl implements OnnxOpaque {
@@ -18,10 +19,10 @@ final class OnnxOpaqueImpl extends OnnxValueImpl implements OnnxOpaque {
         this.opaqueInfo = opaqueInfo;
     }
 
-    //    @Override
-    //    public OnnxOpaque asOpaque() {
-    //        return this;
-    //    }
+    // @Override
+    // public OnnxOpaque asOpaque() {
+    // return this;
+    // }
 
     @Override
     public OpaqueInfo getInfo() {
@@ -35,13 +36,18 @@ public String toString() {
 
     @Override
     public MemoryAddress toNative(
-            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, MemorySession allocator) {
+            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, SegmentAllocator allocator) {
         // TODO Auto-generated method stub
         return null;
     }
 
     @Override
-    public void fromNative(ApiImpl api, MemoryAddress ortAllocator, MemoryAddress address, MemorySession allocator) {
+    public void fromNative(
+            ApiImpl api,
+            MemoryAddress ortAllocator,
+            MemoryAddress address,
+            SegmentAllocator allocator,
+            MemorySession session) {
         // TODO Auto-generated method stub
 
     }

diff --git a/src/main/java/com/jyuzawa/onnxruntime/OnnxOptionalImpl.java b/src/main/java/com/jyuzawa/onnxruntime/OnnxOptionalImpl.java
@@ -6,6 +6,7 @@
 
 import java.lang.foreign.MemoryAddress;
 import java.lang.foreign.MemorySession;
+import java.lang.foreign.SegmentAllocator;
 import java.util.NoSuchElementException;
 
 final class OnnxOptionalImpl extends OnnxValueImpl implements OnnxOptional {
@@ -55,13 +56,18 @@ public OnnxValue set() {
 
     @Override
     public MemoryAddress toNative(
-            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, MemorySession allocator) {
+            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, SegmentAllocator allocator) {
         // TODO Auto-generated method stub
         return null;
     }
 
     @Override
-    public void fromNative(ApiImpl api, MemoryAddress ortAllocator, MemoryAddress address, MemorySession allocator) {
+    public void fromNative(
+            ApiImpl api,
+            MemoryAddress ortAllocator,
+            MemoryAddress address,
+            SegmentAllocator allocator,
+            MemorySession session) {
         // TODO Auto-generated method stub
 
     }

diff --git a/src/main/java/com/jyuzawa/onnxruntime/OnnxSequenceImpl.java b/src/main/java/com/jyuzawa/onnxruntime/OnnxSequenceImpl.java
@@ -9,6 +9,7 @@
 import java.lang.foreign.MemoryAddress;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.MemorySession;
+import java.lang.foreign.SegmentAllocator;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -57,36 +58,34 @@ public OnnxValue add() {
 
     @Override
     public MemoryAddress toNative(
-            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, MemorySession allocator) {
+            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, SegmentAllocator allocator) {
         int size = data.size();
         MemorySegment valuesArray = allocator.allocateArray(C_POINTER, size);
         for (int i = 0; i < size; i++) {
             OnnxValueImpl value = data.get(i);
             valuesArray.setAtIndex(C_POINTER, i, value.toNative(api, ortAllocator, memoryInfo, allocator));
         }
-        MemoryAddress value = api.create(
+        return api.create(
                 allocator,
                 out -> api.CreateValue.apply(valuesArray.address(), size, OnnxType.SEQUENCE.getNumber(), out));
-        allocator.addCloseAction(() -> {
-            api.ReleaseValue.apply(value);
-        });
-        return value;
     }
 
     @Override
-    public void fromNative(ApiImpl api, MemoryAddress ortAllocator, MemoryAddress address, MemorySession allocator) {
+    public void fromNative(
+            ApiImpl api,
+            MemoryAddress ortAllocator,
+            MemoryAddress address,
+            SegmentAllocator allocator,
+            MemorySession session) {
         long outputs = api.extractLong(allocator, out -> api.GetValueCount.apply(address, out));
         for (int i = 0; i < outputs; i++) {
             final int index = i;
             MemoryAddress valueAddress =
                     api.create(allocator, out -> api.GetValue.apply(address, index, ortAllocator, out));
             OnnxValueImpl value = OnnxValueImpl.fromTypeInfo(typeInfo);
-            value.fromNative(api, ortAllocator, valueAddress, allocator);
+            value.fromNative(api, ortAllocator, valueAddress, allocator, session);
             data.add(value);
         }
-        allocator.addCloseAction(() -> {
-            api.ReleaseValue.apply(address);
-        });
     }
 
     @Override

diff --git a/src/main/java/com/jyuzawa/onnxruntime/OnnxTensorBufferImpl.java b/src/main/java/com/jyuzawa/onnxruntime/OnnxTensorBufferImpl.java
@@ -7,6 +7,7 @@
 import java.lang.foreign.MemoryAddress;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.MemorySession;
+import java.lang.foreign.SegmentAllocator;
 import java.nio.Buffer;
 import java.util.function.IntFunction;
 
@@ -26,13 +27,13 @@ public final String toString() {
 
     @Override
     public final MemoryAddress toNative(
-            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, MemorySession allocator) {
+            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, SegmentAllocator allocator) {
         MemorySegment rawInputData = getMemorySegment();
         // TODO: move value layout to this class?
         MemorySegment inputData =
                 allocator.allocateArray(tensorInfo.getType().getValueLayout(), rawInputData.byteSize());
         inputData.copyFrom(rawInputData);
-        MemoryAddress tensor = api.create(
+        return api.create(
                 allocator,
                 out -> api.CreateTensorWithDataAsOrtValue.apply(
                         memoryInfo,
@@ -42,21 +43,18 @@ public final MemoryAddress toNative(
                         tensorInfo.getShape().size(),
                         tensorInfo.getType().getNumber(),
                         out));
-        allocator.addCloseAction(() -> {
-            api.ReleaseValue.apply(tensor);
-        });
-        return tensor;
     }
 
     @Override
     public final void fromNative(
-            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress address, MemorySession allocator) {
+            ApiImpl api,
+            MemoryAddress ortAllocator,
+            MemoryAddress address,
+            SegmentAllocator allocator,
+            MemorySession session) {
         MemoryAddress floatOutput = api.create(allocator, out -> api.GetTensorMutableData.apply(address, out));
-        MemorySegment segment = MemorySegment.ofAddress(floatOutput, tensorInfo.getByteCount(), allocator);
+        MemorySegment segment = MemorySegment.ofAddress(floatOutput, tensorInfo.getByteCount(), session);
         getMemorySegment().copyFrom(segment);
-        allocator.addCloseAction(() -> {
-            api.ReleaseValue.apply(address);
-        });
     }
 
     protected abstract MemorySegment getMemorySegment();

diff --git a/src/main/java/com/jyuzawa/onnxruntime/OnnxTensorStringImpl.java b/src/main/java/com/jyuzawa/onnxruntime/OnnxTensorStringImpl.java
@@ -10,6 +10,7 @@
 import java.lang.foreign.MemoryAddress;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.MemorySession;
+import java.lang.foreign.SegmentAllocator;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Iterator;
@@ -36,8 +37,7 @@ public String[] getStringBuffer() {
 
     @Override
     public MemoryAddress toNative(
-            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, MemorySession allocator) {
-
+            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, SegmentAllocator allocator) {
         int numOutputs = buffer.length;
         MemorySegment stringArray = allocator.allocateArray(C_POINTER, numOutputs);
         for (int i = 0; i < numOutputs; i++) {
@@ -52,14 +52,16 @@ public MemoryAddress toNative(
                         tensorInfo.getType().getNumber(),
                         out));
         api.checkStatus(api.FillStringTensor.apply(tensor, stringArray.address(), numOutputs));
-        allocator.addCloseAction(() -> {
-            api.ReleaseValue.apply(tensor);
-        });
         return tensor;
     }
 
     @Override
-    public void fromNative(ApiImpl api, MemoryAddress ortAllocator, MemoryAddress address, MemorySession allocator) {
+    public void fromNative(
+            ApiImpl api,
+            MemoryAddress ortAllocator,
+            MemoryAddress address,
+            SegmentAllocator allocator,
+            MemorySession session) {
         int numOutputs = buffer.length;
         for (int i = 0; i < numOutputs; i++) {
             final long index = i;
@@ -70,9 +72,6 @@ public void fromNative(ApiImpl api, MemoryAddress ortAllocator, MemoryAddress ad
             api.checkStatus(api.GetStringTensorElement.apply(address, length, index, output.address()));
             buffer[i] = output.getUtf8String(0);
         }
-        allocator.addCloseAction(() -> {
-            api.ReleaseValue.apply(address);
-        });
     }
 
     @Override

diff --git a/src/main/java/com/jyuzawa/onnxruntime/OnnxValueImpl.java b/src/main/java/com/jyuzawa/onnxruntime/OnnxValueImpl.java
@@ -6,6 +6,7 @@
 
 import java.lang.foreign.MemoryAddress;
 import java.lang.foreign.MemorySession;
+import java.lang.foreign.SegmentAllocator;
 import java.util.NoSuchElementException;
 
 abstract class OnnxValueImpl implements OnnxValue {
@@ -52,9 +53,14 @@ public OnnxMap asMap() {
     // }
 
     abstract MemoryAddress toNative(
-            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, MemorySession scope);
+            ApiImpl api, MemoryAddress ortAllocator, MemoryAddress memoryInfo, SegmentAllocator allocator);
 
-    abstract void fromNative(ApiImpl api, MemoryAddress ortAllocator, MemoryAddress address, MemorySession scope);
+    abstract void fromNative(
+            ApiImpl api,
+            MemoryAddress ortAllocator,
+            MemoryAddress address,
+            SegmentAllocator allocator,
+            MemorySession session);
 
     static final OnnxValueImpl fromTypeInfo(TypeInfoImpl typeInfo) {
         OnnxType type = typeInfo.getType();

diff --git a/src/main/java/com/jyuzawa/onnxruntime/TensorInfoImpl.java b/src/main/java/com/jyuzawa/onnxruntime/TensorInfoImpl.java
@@ -7,7 +7,7 @@
 import static com.jyuzawa.onnxruntime_extern.onnxruntime_all_h.C_LONG;
 
 import java.lang.foreign.MemorySegment;
-import java.lang.foreign.MemorySession;
+import java.lang.foreign.SegmentAllocator;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -31,7 +31,7 @@ final class TensorInfoImpl implements TensorInfo {
         this.elementCount = elementCount;
     }
 
-    static TensorInfoImpl of(OnnxTensorElementDataType type, long elementCount, MemorySession scope) {
+    static TensorInfoImpl of(OnnxTensorElementDataType type, long elementCount, SegmentAllocator scope) {
         MemorySegment shapeData = scope.allocateArray(C_LONG, new long[] {elementCount});
         return new TensorInfoImpl(type, shapeData, 1, elementCount);
     }