let paddle hold the creation buffer (#1162)

* let paddle hold the creation buffer * use the data
deepjavalibrary · Aug 11, 2021 · c09b449 · c09b449
1 parent ea814b4
commit c09b449
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpNDArray.java b/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpNDArray.java
@@ -26,18 +26,22 @@
 public class PpNDArray extends NativeResource<Long> implements NDArrayAdapter {
 
     private PpNDManager manager;
+    // we keep the data to prevent GC from early collecting native memory
+    private ByteBuffer data;
     private Shape shape;
     private DataType dataType;
 
     /**
      * Constructs an PpNDArray from a native handle (internal. Use {@link NDManager} instead).
      *
      * @param manager the manager to attach the new array to
+     * @param data bytebuffer that holds the native memory
      * @param handle the pointer to the native MxNDArray memory
      */
-    public PpNDArray(PpNDManager manager, long handle) {
+    public PpNDArray(PpNDManager manager, ByteBuffer data, long handle) {
         super(handle);
         this.manager = manager;
+        this.data = data;
         manager.attachInternal(getUid(), this);
     }
 
@@ -130,7 +134,11 @@ public void detach() {
     /** {@inheritDoc} */
     @Override
     public ByteBuffer toByteBuffer() {
-        return JniUtils.getByteBufferFromNd(this);
+        if (data == null) {
+            data = JniUtils.getByteBufferFromNd(this);
+        }
+        data.rewind();
+        return data;
     }
 
     /** {@inheritDoc} */
@@ -148,6 +156,7 @@ public void close() {
         Long pointer = handle.getAndSet(null);
         if (pointer != null) {
             JniUtils.deleteNd(pointer);
+            data = null;
         }
     }
 }
diff --git a/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/jni/JniUtils.java b/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/jni/JniUtils.java
@@ -37,7 +37,7 @@ public static PpNDArray createNdArray(
         long handle =
                 PaddleLibrary.LIB.paddleCreateTensor(
                         data, data.remaining(), intShape, PpDataType.toPaddlePaddle(dtype));
-        return new PpNDArray(manager, handle);
+        return new PpNDArray(manager, data, handle);
     }
 
     public static DataType getDTypeFromNd(PpNDArray array) {
@@ -119,7 +119,7 @@ public static PpNDArray[] predictorForward(
         PpNDManager manager = (PpNDManager) inputs[0].getManager();
         PpNDArray[] arrays = new PpNDArray[outputs.length];
         for (int i = 0; i < outputs.length; i++) {
-            arrays[i] = new PpNDArray(manager, outputs[i]);
+            arrays[i] = new PpNDArray(manager, null, outputs[i]);
         }
         return arrays;
     }