[android] Add support for float16 tensor (pytorch#15479)

larryliu0820 · web-flow · commit dcc9a9113d80 · 2025-10-31T00:01:23.000-07:00
The Android binding exposes helpers for feeding IEEE-754 half-precision
(FP16) inputs directly.
Use `Tensor.fromBlob(shortArray, shape)` or reuse a direct `ShortBuffer`
created via
`Tensor.allocateHalfBuffer(numElements)` to avoid extra copies:

```kotlin
val shape = longArrayOf(24, 4096)
val halfData: ShortArray = buildHalfEncodedData()
val tensor = Tensor.fromBlob(halfData, shape)

val buffer = Tensor.allocateHalfBuffer(halfData.size)
buffer.put(halfData)
buffer.rewind()
val tensorNoCopy = Tensor.fromBlob(buffer, shape)
```

All buffers must be direct and use the native byte order; the helper
above takes care of this.
diff --git a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/Tensor.java b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/Tensor.java
@@ -18,6 +18,7 @@
 import java.nio.FloatBuffer;
 import java.nio.IntBuffer;
 import java.nio.LongBuffer;
+import java.nio.ShortBuffer;
 import java.util.Arrays;
 import java.util.Locale;
 import org.pytorch.executorch.annotations.Experimental;
@@ -57,6 +58,7 @@ public abstract class Tensor {
   private static final int BYTE_SIZE_BYTES = 1;
   private static final int INT_SIZE_BYTES = 4;
   private static final int LONG_SIZE_BYTES = 8;
+  private static final int HALF_SIZE_BYTES = 2;
   private static final int FLOAT_SIZE_BYTES = 4;
   private static final int DOUBLE_SIZE_BYTES = 8;
 
@@ -107,6 +109,18 @@ public static LongBuffer allocateLongBuffer(int numElements) {
         .asLongBuffer();
   }
 
+  /**
+   * Allocates a new direct {@link ShortBuffer} with native byte order and specified capacity that
+   * can be used in {@link Tensor#fromBlob(ShortBuffer, long[])}.
+   *
+   * @param numElements capacity (number of elements) of result buffer.
+   */
+  public static ShortBuffer allocateHalfBuffer(int numElements) {
+    return ByteBuffer.allocateDirect(numElements * HALF_SIZE_BYTES)
+        .order(ByteOrder.nativeOrder())
+        .asShortBuffer();
+  }
+
   /**
    * Allocates a new direct {@link DoubleBuffer} with native byte order with specified capacity that
    * can be used in {@link Tensor#fromBlob(DoubleBuffer, long[])}.
@@ -187,6 +201,23 @@ public static Tensor fromBlob(float[] data, long[] shape) {
     return new Tensor_float32(floatBuffer, shape);
   }
 
+  /**
+   * Creates a new Tensor instance with dtype torch.float16 with specified shape and data as array
+   * of IEEE-754 half-precision values encoded in {@code short}s.
+   *
+   * @param data Tensor elements encoded as 16-bit floats.
+   * @param shape Tensor shape
+   */
+  public static Tensor fromBlob(short[] data, long[] shape) {
+    checkArgument(data != null, ERROR_MSG_DATA_ARRAY_NOT_NULL);
+    checkArgument(shape != null, ERROR_MSG_SHAPE_NOT_NULL);
+    checkShape(shape);
+    checkShapeAndDataCapacityConsistency(data.length, shape);
+    final ShortBuffer shortBuffer = allocateHalfBuffer((int) numel(shape));
+    shortBuffer.put(data);
+    return new Tensor_float16(shortBuffer, shape);
+  }
+
   /**
    * Creates a new Tensor instance with dtype torch.int64 with specified shape and data as array of
    * longs.
@@ -301,6 +332,26 @@ public static Tensor fromBlob(FloatBuffer data, long[] shape) {
     return new Tensor_float32(data, shape);
   }
 
+  /**
+   * Creates a new Tensor instance with dtype torch.float16 with specified shape and data.
+   *
+   * @param data Direct buffer with native byte order that contains {@code Tensor.numel(shape)}
+   *     elements encoded as IEEE-754 half-precision floats. The buffer is used directly without
+   *     copying.
+   * @param shape Tensor shape
+   */
+  public static Tensor fromBlob(ShortBuffer data, long[] shape) {
+    checkArgument(data != null, ERROR_MSG_DATA_BUFFER_NOT_NULL);
+    checkArgument(shape != null, ERROR_MSG_SHAPE_NOT_NULL);
+    checkShape(shape);
+    checkShapeAndDataCapacityConsistency(data.capacity(), shape);
+    checkArgument(data.isDirect(), ERROR_MSG_DATA_BUFFER_MUST_BE_DIRECT);
+    checkArgument(
+        (data.order() == ByteOrder.nativeOrder()),
+        ERROR_MSG_DATA_BUFFER_MUST_HAVE_NATIVE_BYTE_ORDER);
+    return new Tensor_float16(data, shape);
+  }
+
   /**
    * Creates a new Tensor instance with dtype torch.int64 with specified shape and data.
    *
@@ -388,6 +439,16 @@ public byte[] getDataAsByteArray() {
         "Tensor of type " + getClass().getSimpleName() + " cannot return data as byte array.");
   }
 
+  /**
+   * @return a Java short array that contains the tensor data interpreted as IEEE-754 half-precision
+   *     bit patterns. This may be a copy or reference.
+   * @throws IllegalStateException if it is called for a non-float16 tensor.
+   */
+  public short[] getDataAsShortArray() {
+    throw new IllegalStateException(
+        "Tensor of type " + getClass().getSimpleName() + " cannot return data as short array.");
+  }
+
   /**
    * @return a Java byte array that contains the tensor data. This may be a copy or reference.
    * @throws IllegalStateException if it is called for a non-uint8 tensor.
@@ -569,6 +630,74 @@ public String toString() {
     }
   }
 
+  static class Tensor_float16 extends Tensor {
+    private final ShortBuffer data;
+
+    private Tensor_float16(ShortBuffer data, long[] shape) {
+      super(shape);
+      this.data = data;
+    }
+
+    @Override
+    public DType dtype() {
+      return DType.HALF;
+    }
+
+    @Override
+    Buffer getRawDataBuffer() {
+      return data;
+    }
+
+    @Override
+    public short[] getDataAsShortArray() {
+      data.rewind();
+      short[] arr = new short[data.remaining()];
+      data.get(arr);
+      return arr;
+    }
+
+    @Override
+    public float[] getDataAsFloatArray() {
+      data.rewind();
+      int remaining = data.remaining();
+      float[] arr = new float[remaining];
+      for (int i = 0; i < remaining; i++) {
+        arr[i] = halfBitsToFloat(data.get());
+      }
+      return arr;
+    }
+
+    @Override
+    public String toString() {
+      return String.format("Tensor(%s, dtype=torch.float16)", Arrays.toString(shape));
+    }
+
+    private static float halfBitsToFloat(short halfBits) {
+      int h = halfBits & 0xFFFF;
+      int sign = (h >>> 15) & 0x1;
+      int exp = (h >>> 10) & 0x1F;
+      int mant = h & 0x3FF;
+
+      if (exp == 0) {
+        if (mant == 0) {
+          return sign == 0 ? 0.0f : -0.0f;
+        }
+        float result = mant * 5.9604645e-8f; // 2^-24
+        return sign == 0 ? result : -result;
+      } else if (exp == 0x1F) {
+        if (mant == 0) {
+          return sign == 0 ? Float.POSITIVE_INFINITY : Float.NEGATIVE_INFINITY;
+        }
+        int bits = (sign << 31) | 0x7f800000 | (mant << 13);
+        return Float.intBitsToFloat(bits);
+      } else {
+        int exp32 = exp + 112; // 127 (float bias) - 15 (half bias)
+        int bits = (sign << 31) | (exp32 << 23) | (mant << 13);
+        return Float.intBitsToFloat(bits);
+      }
+    }
+  }
+
   static class Tensor_int64 extends Tensor {
     private final LongBuffer data;
 
@@ -691,6 +820,8 @@ private static Tensor nativeNewTensor(
 
     if (DType.FLOAT.jniCode == dtype) {
       tensor = new Tensor_float32(data.asFloatBuffer(), shape);
+    } else if (DType.HALF.jniCode == dtype) {
+      tensor = new Tensor_float16(data.asShortBuffer(), shape);
     } else if (DType.INT32.jniCode == dtype) {
       tensor = new Tensor_int32(data.asIntBuffer(), shape);
     } else if (DType.INT64.jniCode == dtype) {
@@ -727,6 +858,11 @@ public byte[] toByteArray() {
       tensorAsByteArray = new byte[(int) numel()];
       Tensor_int8 thiz = (Tensor_int8) this;
       ByteBuffer.wrap(tensorAsByteArray).put(thiz.getDataAsByteArray());
+    } else if (dtype() == DType.HALF) {
+      dtypeSize = HALF_SIZE_BYTES;
+      tensorAsByteArray = new byte[(int) numel() * dtypeSize];
+      Tensor_float16 thiz = (Tensor_float16) this;
+      ByteBuffer.wrap(tensorAsByteArray).asShortBuffer().put(thiz.getDataAsShortArray());
     } else if (dtype() == DType.INT16) {
       throw new IllegalArgumentException("DType.INT16 is not supported in Java so far");
     } else if (dtype() == DType.INT32) {
@@ -794,6 +930,8 @@ public static Tensor fromByteArray(byte[] bytes) {
       return new Tensor_uint8(buffer, shape);
     } else if (dtype == DType.INT8.jniCode) {
       return new Tensor_int8(buffer, shape);
+    } else if (dtype == DType.HALF.jniCode) {
+      return new Tensor_float16(buffer.asShortBuffer(), shape);
     } else if (dtype == DType.INT32.jniCode) {
       return new Tensor_int32(buffer.asIntBuffer(), shape);
     } else if (dtype == DType.INT64.jniCode) {
diff --git a/extension/android/executorch_android/src/test/java/org/pytorch/executorch/TensorTest.kt b/extension/android/executorch_android/src/test/java/org/pytorch/executorch/TensorTest.kt
@@ -7,8 +7,11 @@
  */
 package org.pytorch.executorch
 
+import java.nio.ByteOrder
 import org.assertj.core.api.Assertions.assertThatThrownBy
+import org.junit.Assert.assertArrayEquals
 import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
@@ -184,6 +187,65 @@ class TensorTest {
     assertEquals(data[3].toLong(), tensor.dataAsUnsignedByteArray[3].toLong())
   }
 
+  @Test
+  fun testHalfTensorFromShortArrayAndBuffer() {
+    val data =
+        shortArrayOf(
+            0x3C00.toShort(), // 1.0
+            0xC000.toShort(), // -2.0
+            0x0000.toShort(), // 0.0
+            0x7C00.toShort(), // +inf
+        )
+    val shape = longArrayOf(2, 2)
+    var tensor = Tensor.fromBlob(data, shape)
+    assertEquals(DType.HALF, tensor.dtype())
+    assertEquals(shape[0], tensor.shape()[0])
+    assertEquals(shape[1], tensor.shape()[1])
+    assertEquals(4, tensor.numel())
+    assertArrayEquals(data, tensor.dataAsShortArray)
+    val floats = tensor.dataAsFloatArray
+    assertEquals(1.0f.toDouble(), floats[0].toDouble(), 1e-6)
+    assertEquals((-2.0f).toDouble(), floats[1].toDouble(), 1e-6)
+    assertEquals(0.0f.toDouble(), floats[2].toDouble(), 1e-6)
+    assertEquals(Float.POSITIVE_INFINITY.toDouble(), floats[3].toDouble(), 0.0)
+
+    val buffer = Tensor.allocateHalfBuffer(data.size)
+    assertTrue(buffer.isDirect)
+    assertEquals(ByteOrder.nativeOrder(), buffer.order())
+    buffer.put(data)
+    buffer.rewind()
+
+    tensor = Tensor.fromBlob(buffer, longArrayOf(data.size.toLong()))
+    assertEquals(DType.HALF, tensor.dtype())
+    assertEquals(data.size.toLong(), tensor.shape()[0])
+    assertEquals(data.size.toLong(), tensor.numel())
+    assertArrayEquals(data, tensor.dataAsShortArray)
+    val raw = tensor.rawDataBuffer as java.nio.ShortBuffer
+    assertTrue(raw === buffer)
+  }
+
+  @Test
+  fun testHalfTensorSerializationRoundTrip() {
+    val data =
+        shortArrayOf(
+            0x0000.toShort(),
+            0x0400.toShort(),
+            0x3C00.toShort(),
+            0x7BFF.toShort(),
+        )
+    val shape = longArrayOf(2, 2)
+    val tensor = Tensor.fromBlob(data, shape)
+    val serialized = tensor.toByteArray()
+    val deserialized = Tensor.fromByteArray(serialized)
+
+    assertEquals(DType.HALF, deserialized.dtype())
+    assertEquals(shape[0], deserialized.shape()[0])
+    assertEquals(shape[1], deserialized.shape()[1])
+    assertEquals(4, deserialized.numel())
+    assertArrayEquals(data, deserialized.dataAsShortArray)
+    assertEquals(1.0f.toDouble(), deserialized.dataAsFloatArray[2].toDouble(), 1e-6)
+  }
+
   @Test
   fun testIllegalDataTypeException() {
     val data = floatArrayOf(Float.MIN_VALUE, 0f, 0.1f, Float.MAX_VALUE)
diff --git a/extension/android/jni/jni_layer.cpp b/extension/android/jni/jni_layer.cpp
@@ -13,6 +13,7 @@
 #include <executorch/extension/module/module.h>
 #include <executorch/extension/runner_util/inputs.h>
 #include <executorch/extension/tensor/tensor.h>
+#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/core/portable_type/tensor_impl.h>
 #include <executorch/runtime/platform/log.h>
 #include <executorch/runtime/platform/platform.h>
@@ -117,7 +118,7 @@ class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
     std::vector<executorch::aten::SizesType> shape_vec;
     shape_vec.reserve(rank);
 
-    auto numel = 1;
+    int64_t numel = 1;
     for (int i = 0; i < rank; i++) {
       shape_vec.push_back(shapeArr[i]);
     }
@@ -132,11 +133,24 @@ class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
           static_cast<uint32_t>(Error::InvalidArgument), ss.str().c_str());
     }
     ScalarType scalar_type = java_dtype_to_scalar_type.at(jdtype);
-    const auto dataCapacity = jni->GetDirectBufferCapacity(jbuffer.get());
-    if (dataCapacity != numel) {
+    const jlong dataCapacity = jni->GetDirectBufferCapacity(jbuffer.get());
+    if (dataCapacity < 0) {
+      std::stringstream ss;
+      ss << "Tensor buffer is not direct or has invalid capacity";
+      jni_helper::throwExecutorchException(
+          static_cast<uint32_t>(Error::InvalidArgument), ss.str().c_str());
+    }
+    const size_t elementSize = executorch::runtime::elementSize(scalar_type);
+    const jlong expectedElements = static_cast<jlong>(numel);
+    const jlong expectedBytes =
+        expectedElements * static_cast<jlong>(elementSize);
+    const bool matchesElements = dataCapacity == expectedElements;
+    const bool matchesBytes = dataCapacity == expectedBytes;
+    if (!matchesElements && !matchesBytes) {
       std::stringstream ss;
       ss << "Tensor dimensions(elements number: " << numel
-         << "inconsistent with buffer capacity " << dataCapacity << "]";
+         << ") inconsistent with buffer capacity " << dataCapacity
+         << " (element size bytes: " << elementSize << ")";
       jni_helper::throwExecutorchException(
           static_cast<uint32_t>(Error::InvalidArgument), ss.str().c_str());
     }
