Java Tensor and EValue serialization

WIP

Author: kirklandsign

extension/android/build.gradle

@@ -22,3 +22,7 @@ task makeJar(type: Jar) {
         implementation 'com.facebook.soloader:nativeloader:0.10.5'
     }
 }
+
+dependencies {
+    testImplementation "junit:junit:4.12"
+}

extension/android/src/main/java/org/pytorch/executorch/EValue.java

@@ -9,6 +9,8 @@
 package org.pytorch.executorch;
 
 import com.facebook.jni.annotations.DoNotStrip;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
 import java.util.Locale;
 import java.util.Optional;
 import org.pytorch.executorch.annotations.Experimental;
@@ -287,4 +289,64 @@ private void preconditionType(int typeCodeExpected, int typeCode) {
   private String getTypeName(int typeCode) {
     return typeCode >= 0 && typeCode < TYPE_NAMES.length ? TYPE_NAMES[typeCode] : "Unknown";
   }
+
+  /**
+   * Serializes an {@code EValue} into a {@link ByteBuffer}.
+   * @return The serialized {@code ByteBuffer}.
+   * 
+   * @apiNote This method is experimental and subject to change without notice.
+   * This does NOT supoprt list type.
+   */
+  public ByteBuffer toByteBuffer() {
+    if (isNone()) {
+      return ByteBuffer.allocateDirect(1).put((byte) TYPE_CODE_NONE);
+    } else if (isTensor()) {
+      Tensor t = toTensor();
+      ByteBuffer tByteBuffer = t.toByteBuffer();
+      return ByteBuffer.allocateDirect(1 + tByteBuffer.array().length).put((byte) TYPE_CODE_TENSOR).put(tByteBuffer);
+    } else if (isBool()) {
+      return ByteBuffer.allocateDirect(2).put((byte) TYPE_CODE_BOOL).put((byte) (toBool() ? 1 : 0));
+    } else if (isInt()) {
+      return ByteBuffer.allocateDirect(9).put((byte) TYPE_CODE_INT).putLong(toInt());
+    } else if (isDouble()) {
+      return ByteBuffer.allocateDirect(9).put((byte) TYPE_CODE_DOUBLE).putDouble(toDouble());
+    } else if (isString()) {
+      return ByteBuffer.allocateDirect(1 + toString().length()).put((byte) TYPE_CODE_STRING).put(toString().getBytes());
+    } else {
+      throw new IllegalArgumentException("Unknown Tensor dtype");
+    }
+  }
+
+  /**
+   * Deserializes an {@code EValue} from a {@link ByteBuffer}.
+   * @param buffer The {@link ByteBuffer} to deserialize from.
+   * @return The deserialized {@code EValue}.
+   * 
+   * @apiNote This method is experimental and subject to change without notice.
+   * This does NOT list type.
+   */
+  public static EValue fromByteBuffer(ByteBuffer buffer) {
+    if (buffer == null) {
+      throw new IllegalArgumentException("buffer cannot be null");
+    }
+    if (!buffer.hasRemaining()) {
+      throw new IllegalArgumentException("invalid buffer");
+    }
+    int typeCode = buffer.get();
+    switch (typeCode) {
+      case TYPE_CODE_NONE:
+        return new EValue(TYPE_CODE_NONE);
+      case TYPE_CODE_TENSOR:
+        return from(Tensor.fromByteBuffer(buffer));
+      case TYPE_CODE_STRING:
+        return from(StandardCharsets.UTF_8.decode(buffer).toString());
+      case TYPE_CODE_DOUBLE:
+        return from(buffer.getDouble());
+      case TYPE_CODE_INT:
+        return from(buffer.getLong());
+      case TYPE_CODE_BOOL:
+        return from(buffer.get() != 0);
+    }
+    throw new IllegalArgumentException("invalid type code: " + typeCode);
+  }
 }

extension/android/src/main/java/org/pytorch/executorch/Tensor.java

@@ -10,6 +10,8 @@
 
 import com.facebook.jni.HybridData;
 import com.facebook.jni.annotations.DoNotStrip;
+import java.io.ByteArrayOutputStream;
+import java.io.ObjectOutputStream;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -679,4 +681,99 @@ private static Tensor nativeNewTensor(
     tensor.mHybridData = hybridData;
     return tensor;
   }
+
+  private static final byte[] intToByteArray(int value) {
+    return new byte[] {
+            (byte)(value >>> 24),
+            (byte)(value >>> 16),
+            (byte)(value >>> 8),
+            (byte)value};
+  }
+
+  /**
+   * Serializes a {@code Tensor} into a {@link ByteBuffer}.
+   * @return The serialized {@code ByteBuffer}.
+   * 
+   * @apiNote This method is experimental and subject to change without notice.
+   * This does NOT supoprt list type.
+   */
+  public ByteBuffer toByteBuffer() {
+    int dtypeSize = 0;
+    if (dtype() == DType.FLOAT) {
+      dtypeSize = 4;
+    } else if (dtype() == DType.DOUBLE) {
+      dtypeSize = 8;
+    } else if (dtype() == DType.UINT8) {
+      dtypeSize = 1;
+    } else if (dtype() == DType.INT8) {
+      dtypeSize = 1;
+    } else if (dtype() == DType.INT16) {
+      dtypeSize = 2;
+    } else if (dtype() == DType.INT32) {
+      dtypeSize = 4;
+    } else if (dtype() == DType.INT64) {
+      dtypeSize = 8;
+    } else {
+      throw new IllegalArgumentException("Unknown Tensor dtype");
+    }
+    ByteBuffer byteBuffer = ByteBuffer.allocateDirect(1 + 1 + 4 * shape.length + dtypeSize * (int) numel());
+    byteBuffer.put((byte) dtype().jniCode);
+    byteBuffer.put((byte) shape.length);
+    for (long s : shape) {
+      byteBuffer.put(intToByteArray((int) s));
+    }
+    ByteArrayOutputStream bos = new ByteArrayOutputStream();
+    try (ObjectOutputStream out = new ObjectOutputStream(bos)) {
+      out.writeObject(getRawDataBuffer());
+      out.flush();
+      byteBuffer.put(bos.toByteArray());
+    } catch (Exception ex) {
+      throw new RuntimeException(ex);
+    }
+    return byteBuffer;
+  }
+
+  /**
+   * Deserializes a {@code Tensor} from a {@link ByteBuffer}.
+   * @param buffer The {@link ByteBuffer} to deserialize from.
+   * @return The deserialized {@code Tensor}.
+   * 
+   * @apiNote This method is experimental and subject to change without notice.
+   * This does NOT supoprt list type.
+   */
+  public static Tensor fromByteBuffer(ByteBuffer buffer) {
+    if (buffer == null) {
+      throw new IllegalArgumentException("buffer cannot be null");
+    }
+    if (!buffer.hasRemaining()) {
+      throw new IllegalArgumentException("invalid buffer");
+    }
+    byte scalarType = buffer.get();
+    byte numberOfDimensions = buffer.get();
+    long[] shape = new long[(int) numberOfDimensions];
+    long numel = 1;
+    for (int i = 0; i < numberOfDimensions; i++) {
+      int dim = buffer.getInt();
+      if (dim < 0) {
+        throw new IllegalArgumentException("invalid shape");
+      }
+      shape[i] = dim;
+      numel *= dim;
+    }
+    if (scalarType == DType.FLOAT.jniCode) {
+      return new Tensor_float32(buffer.asFloatBuffer(), shape);
+    } else if (scalarType == DType.DOUBLE.jniCode) {
+      return new Tensor_float64(buffer.asDoubleBuffer(), shape);
+    } else if (scalarType == DType.UINT8.jniCode) {
+      return new Tensor_uint8(buffer, shape);
+    } else if (scalarType == DType.INT8.jniCode) {
+      return new Tensor_int8(buffer, shape);
+    } else if (scalarType == DType.INT16.jniCode) {
+      return new Tensor_int32(buffer.asIntBuffer(), shape);
+    } else if (scalarType == DType.INT64.jniCode) {
+      return new Tensor_int64(buffer.asLongBuffer(), shape);
+    } else {
+      throw new IllegalArgumentException("Unknown Tensor dtype");
+    }
+  }
 }

extension/android/src/test/java/org/pytorch/executorch/TensorTest.java

@@ -0,0 +1,23 @@
+package org.pytorch.executorch;
+
+import java.nio.ByteBuffer;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertEquals;
+public class TensorTest  {
+  @Test
+  public void testHello() {
+    EValue evalue = EValue.from(1);
+    ByteBuffer bb = evalue.toByteBuffer();
+    assertEquals(4, bb.get());
+    assertEquals(0, bb.get());
+    assertEquals(0, bb.get());
+    assertEquals(0, bb.get());
+    assertEquals(1, bb.get());
+
+    // Tensor tensor = Tensor.fromBlob(new float[] {1.0f}, new long[] {1});
+    // tensor.toByteBuffer();
+  }
+}