Add Tensor.copyDataInto for allocation-free output reads on Android

[jgibson2]

Summary

Adds Tensor.copyDataInto(<TypedBuffer>) to the Android Java API for every dtype, mirroring the version up for review on pytorch/executorch#19171. Eliminates the per-call float[] / int[] / etc. allocation that getDataAs*Array() performs by letting the caller supply (and reuse) a destination buffer.

Motivation

Profiling on Android (Perfetto) showed output.toTensor().dataAsFloatArray driving substantial ART GC pressure in the depth-inference loop. Each call allocates a fresh Java float[] sized to the tensor's element count and bulk-copies from the underlying off-heap buffer into it. For a model with two [1, 1, 144, 192] fp32 outputs that's ~110 KB × 2 = 220 KB of Java-heap churn per inference. At 2 fps that's 440 KB/sec just from reading outputs — visibly enough to add several young-generation GCs per second, and ~14% of the inference wall time was being lost to run-queue contention as GC suspended app threads.

The native side already exposes the underlying off-heap FloatBuffer as a zero-copy view of the C++ tensor's data_ptr() (via the package-private getRawDataBuffer). The new copyDataInto API gives external callers a public way to drain that buffer into a destination they own and reuse across calls.

API

public void copyDataInto(FloatBuffer dst)        // float32 + float16 (widening)
public void copyDataInto(IntBuffer dst)          // int32
public void copyDataInto(LongBuffer dst)         // int64
public void copyDataInto(DoubleBuffer dst)       // float64
public void copyDataInto(ShortBuffer dst)        // float16 raw bits
public void copyDataInto(ByteBuffer dst)         // int8
public void copyDataIntoUnsigned(ByteBuffer dst) // uint8

• Float16 has both a FloatBuffer overload (per-element half→float widening) and a ShortBuffer overload (raw fp16 bits, no widening).
• int8 / uint8 are split into separate methods to mirror the existing getDataAsByteArray vs getDataAsUnsignedByteArray distinction — calling the wrong one throws.
• All overloads write at the destination's current position and advance it (standard Buffer.put semantics).
• The fp16 → fp32 widening path includes an upfront dst.remaining() capacity check so an undersized destination throws BufferOverflowException before any partial widening is observed (matching the all-or-nothing semantics of the bulk-put paths).

Caller-side usage

// One-time setup
FloatBuffer depthBuf = Tensor.allocateFloatBuffer(numelDepth);

// Per inference
EValue[] outputs = module.forward(...);
depthBuf.rewind();
outputs[0].toTensor().copyDataInto(depthBuf);   // no allocation
// ... read from depthBuf ...

Test plan

• TensorTest.kt covers each variant's happy path, asymmetric int8/uint8 rejection, fp16 raw-bits vs widened paths, position-respecting writes, and BufferOverflowException on undersized destinations (including the fp16 widening path's upfront check).
• Tensor.java compiles standalone with javac 17.
• google-java-format (1.23.0, project standard) and ktfmt (Meta style, matches spotless { kotlin { ktfmt() } } config) both clean.
• Verify against polycam EstimateDepth.kt (forthcoming change to swap dataAsFloatArray for copyDataInto + a 2-slot encoder pool).

🤖 Generated with Claude Code

[meredithbayne]

Nice improvements!