Update Android API

Author: kirklandsign

extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java

@@ -173,23 +173,23 @@ public native int generate(
    * @param height Input image height
    * @param channels Input image number of channels
    * @param startPos The starting position in KV cache of the input in the LLM.
-   * @return The updated starting position in KV cache of the input in the LLM.
+   * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer
+   * exposed to user.
    * @throws RuntimeException if the prefill failed
    */
   @Deprecated
   public long prefillImages(int[] image, int width, int height, int channels, long startPos) {
     if (startPos == 0) {
       resetContext();
     }
-    long[] nativeResult = prefillImagesNative(image, width, height, channels);
-    if (nativeResult[0] != 0) {
+    int nativeResult = prefillImagesNative(image, width, height, channels);
+    if (nativeResult != 0) {
       throw new RuntimeException("Prefill failed with error code: " + nativeResult[0]);
     }
-    return nativeResult[1];
+    return 0;
   }
 
-  // returns a tuple of (status, updated startPos)
-  private native long[] prefillImagesNative(int[] image, int width, int height, int channels);
+  private native int prefillImagesNative(int[] image, int width, int height, int channels);
 
   /**
    * Prefill an LLaVA Module with the given text input.
@@ -199,23 +199,24 @@ public long prefillImages(int[] image, int width, int height, int channels, long
    *     reference and will be updated inside this function.
    * @param bos The number of BOS (begin of sequence) token.
    * @param eos The number of EOS (end of sequence) token.
-   * @return The updated starting position in KV cache of the input in the LLM.
+   * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer
+   * exposed to user.
    * @throws RuntimeException if the prefill failed
    */
   @Deprecated
   public long prefillPrompt(String prompt, long startPos, int bos, int eos) {
     if (startPos == 0) {
       resetContext();
     }
-    long[] nativeResult = prefillPromptNative(prompt, bos, eos);
-    if (nativeResult[0] != 0) {
+    int nativeResult = prefillPromptNative(prompt, bos, eos);
+    if (nativeResult != 0) {
       throw new RuntimeException("Prefill failed with error code: " + nativeResult[0]);
     }
-    return nativeResult[1];
+    return 0;
   }
 
   // returns a tuple of (status, updated startPos)
-  private native long[] prefillPromptNative(String prompt, int bos, int eos);
+  private native int prefillPromptNative(String prompt, int bos, int eos);
 
   /**
    * Generate tokens from the given prompt, starting from the given position.
@@ -238,7 +239,7 @@ public native int generateFromPos(
    *
    * <p>The startPos will be reset to 0.
    */
-  public native int resetContext();
+  public native void resetContext();
 
   /** Stop current generate() before it finishes. */
   @DoNotStrip

extension/android/jni/jni_layer_llama.cpp

@@ -251,42 +251,24 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
     return 0;
   }
 
-  // Returns a tuple of (error, start_pos)
-  // Contract is valid within an AAR (JNI + corresponding Java code)
-  // If the first element is not Error::Ok, the other element is undefined.
-  facebook::jni::local_ref<jlongArray>
+  jint
   prefill_prompt(facebook::jni::alias_ref<jstring> prompt, jint bos, jint eos) {
-    facebook::jni::local_ref<jlongArray> tuple_result =
-        facebook::jni::make_long_array(2);
     if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL) {
-      tuple_result->pin()[0] = static_cast<jint>(Error::NotSupported);
-      return tuple_result;
+      return static_cast<jint>(Error::NotSupported);
     }
 
     auto&& result =
         multi_modal_runner_->prefill_prompt(prompt->toStdString(), bos, eos);
-    tuple_result->pin()[0] = static_cast<jint>(Error::Ok);
-    if (result.ok()) {
-      tuple_result->pin()[1] = static_cast<jlong>(start_pos);
-    }
-    return tuple_result;
+    return static_cast<jint>(result.error());
   }
 
-  // Returns a tuple of (error, start_pos)
-  // Contract is valid within an AAR (JNI + corresponding Java code)
-  // If the first element is not Error::Ok, the other element is undefined.
-
-  facebook::jni::local_ref<jlongArray> prefill_images(
+  jint prefill_images(
       facebook::jni::alias_ref<jintArray> image,
       jint width,
       jint height,
       jint channels) {
-    facebook::jni::local_ref<jlongArray> tuple_result =
-        facebook::jni::make_long_array(2);
-
     if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL) {
-      tuple_result->pin()[0] = static_cast<jint>(Error::NotSupported);
-      return tuple_result;
+      return static_cast<jint>(Error::NotSupported);
     }
 
     auto image_size = image->size();
@@ -301,12 +283,9 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       llm::Image image_runner{image_data, width, height, channels};
       images.push_back(image_runner);
     }
-    // TODO(hsz): make  start_pos a reference and update it here
     jint result =
         static_cast<jint>(multi_modal_runner_->prefill_images(images));
-    tuple_result->pin()[0] = result;
-    tuple_result->pin()[1] = static_cast<jlong>(start_pos);
-    return tuple_result;
+    return result;
   }
 
   jint generate_from_pos(