From c895909f47595b4f544acd67fd18b963d9d4c872 Mon Sep 17 00:00:00 2001
From: Ewan <166796318+ewan0x79@users.noreply.github.com>
Date: Mon, 1 Jul 2024 23:47:37 +0800
Subject: [PATCH] [enhancement] Optimize memory copy overhead to enhance
 performance. (#3289)

---
 api/src/main/java/ai/djl/ndarray/NDManager.java | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/api/src/main/java/ai/djl/ndarray/NDManager.java b/api/src/main/java/ai/djl/ndarray/NDManager.java
index 4cc020ea028..8f26e2713d1 100644
--- a/api/src/main/java/ai/djl/ndarray/NDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/NDManager.java
@@ -399,7 +399,7 @@ default NDArray create(boolean[] data) {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(float[][] data) {
-        FloatBuffer buffer = FloatBuffer.allocate(data.length * data[0].length);
+        FloatBuffer buffer = allocateDirect(data.length * data[0].length * 4).asFloatBuffer();
         for (float[] d : data) {
             buffer.put(d);
         }
@@ -414,7 +414,7 @@ default NDArray create(float[][] data) {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(int[][] data) {
-        IntBuffer buffer = IntBuffer.allocate(data.length * data[0].length);
+        IntBuffer buffer = allocateDirect(data.length * data[0].length * 4).asIntBuffer();
         for (int[] d : data) {
             buffer.put(d);
         }
@@ -429,7 +429,7 @@ default NDArray create(int[][] data) {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(double[][] data) {
-        DoubleBuffer buffer = DoubleBuffer.allocate(data.length * data[0].length);
+        DoubleBuffer buffer = allocateDirect(data.length * data[0].length * 8).asDoubleBuffer();
         for (double[] d : data) {
             buffer.put(d);
         }
@@ -444,7 +444,7 @@ default NDArray create(double[][] data) {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(long[][] data) {
-        LongBuffer buffer = LongBuffer.allocate(data.length * data[0].length);
+        LongBuffer buffer = allocateDirect(data.length * data[0].length * 8).asLongBuffer();
         for (long[] d : data) {
             buffer.put(d);
         }
@@ -459,7 +459,7 @@ default NDArray create(long[][] data) {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(byte[][] data) {
-        ByteBuffer buffer = ByteBuffer.allocate(data.length * data[0].length);
+        ByteBuffer buffer = allocateDirect(data.length * data[0].length);
         for (byte[] d : data) {
             buffer.put(d);
         }
@@ -474,7 +474,7 @@ default NDArray create(byte[][] data) {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(boolean[][] data) {
-        ByteBuffer buffer = ByteBuffer.allocate(data.length * data[0].length);
+        ByteBuffer buffer = allocateDirect(data.length * data[0].length);
         for (boolean[] d : data) {
             for (boolean b : d) {
                 buffer.put((byte) (b ? 1 : 0));