Unify Android JNI to single IRunner, wire prefill to runner

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

@@ -372,136 +372,131 @@ public int generate(
   }
 
   /**
-   * Prefill a multimodal Module with the given images input.
+   * Prefill the KV cache with the given image input.
    *
    * @param image Input image as a byte array
    * @param width Input image width
    * @param height Input image height
    * @param channels Input image number of channels
-   * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer
-   *     exposed to user.
+   * @return 0 on success
    * @throws RuntimeException if the prefill failed
    */
   @Experimental
   public long prefillImages(int[] image, int width, int height, int channels) {
-    int nativeResult = appendImagesInput(image, width, height, channels);
+    int nativeResult = prefillImagesInput(image, width, height, channels);
     if (nativeResult != 0) {
       throw new RuntimeException("Prefill failed with error code: " + nativeResult);
     }
     return 0;
   }
 
-  private native int appendImagesInput(int[] image, int width, int height, int channels);
+  private native int prefillImagesInput(int[] image, int width, int height, int channels);
 
   /**
-   * Prefill a multimodal Module with the given images input.
+   * Prefill the KV cache with the given normalized image input.
    *
    * @param image Input normalized image as a float array
    * @param width Input image width
    * @param height Input image height
    * @param channels Input image number of channels
-   * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer
-   *     exposed to user.
+   * @return 0 on success
    * @throws RuntimeException if the prefill failed
    */
   @Experimental
   public long prefillImages(float[] image, int width, int height, int channels) {
-    int nativeResult = appendNormalizedImagesInput(image, width, height, channels);
+    int nativeResult = prefillNormalizedImagesInput(image, width, height, channels);
     if (nativeResult != 0) {
       throw new RuntimeException("Prefill failed with error code: " + nativeResult);
     }
     return 0;
   }
 
-  private native int appendNormalizedImagesInput(
+  private native int prefillNormalizedImagesInput(
       float[] image, int width, int height, int channels);
 
   /**
-   * Prefill a multimodal Module with the given audio input.
+   * Prefill the KV cache with the given preprocessed audio input.
    *
    * @param audio Input preprocessed audio as a byte array
    * @param batch_size Input batch size
    * @param n_bins Input number of bins
    * @param n_frames Input number of frames
-   * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer
-   *     exposed to user.
+   * @return 0 on success
    * @throws RuntimeException if the prefill failed
    */
   @Experimental
   public long prefillAudio(byte[] audio, int batch_size, int n_bins, int n_frames) {
-    int nativeResult = appendAudioInput(audio, batch_size, n_bins, n_frames);
+    int nativeResult = prefillAudioInput(audio, batch_size, n_bins, n_frames);
     if (nativeResult != 0) {
       throw new RuntimeException("Prefill failed with error code: " + nativeResult);
     }
     return 0;
   }
 
-  private native int appendAudioInput(byte[] audio, int batch_size, int n_bins, int n_frames);
+  private native int prefillAudioInput(byte[] audio, int batch_size, int n_bins, int n_frames);
 
   /**
-   * Prefill a multimodal Module with the given audio input.
+   * Prefill the KV cache with the given preprocessed audio input.
    *
    * @param audio Input preprocessed audio as a float array
    * @param batch_size Input batch size
    * @param n_bins Input number of bins
    * @param n_frames Input number of frames
-   * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer
-   *     exposed to user.
+   * @return 0 on success
    * @throws RuntimeException if the prefill failed
    */
   @Experimental
   public long prefillAudio(float[] audio, int batch_size, int n_bins, int n_frames) {
-    int nativeResult = appendAudioInputFloat(audio, batch_size, n_bins, n_frames);
+    int nativeResult = prefillAudioInputFloat(audio, batch_size, n_bins, n_frames);
     if (nativeResult != 0) {
       throw new RuntimeException("Prefill failed with error code: " + nativeResult);
     }
     return 0;
   }
 
-  private native int appendAudioInputFloat(float[] audio, int batch_size, int n_bins, int n_frames);
+  private native int prefillAudioInputFloat(
+      float[] audio, int batch_size, int n_bins, int n_frames);
 
   /**
-   * Prefill a multimodal Module with the given raw audio input.
+   * Prefill the KV cache with the given raw audio input.
    *
    * @param audio Input raw audio as a byte array
    * @param batch_size Input batch size
    * @param n_channels Input number of channels
    * @param n_samples Input number of samples
-   * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer
-   *     exposed to user.
+   * @return 0 on success
    * @throws RuntimeException if the prefill failed
    */
   @Experimental
   public long prefillRawAudio(byte[] audio, int batch_size, int n_channels, int n_samples) {
-    int nativeResult = appendRawAudioInput(audio, batch_size, n_channels, n_samples);
+    int nativeResult = prefillRawAudioInput(audio, batch_size, n_channels, n_samples);
     if (nativeResult != 0) {
       throw new RuntimeException("Prefill failed with error code: " + nativeResult);
     }
     return 0;
   }
 
-  private native int appendRawAudioInput(
+  private native int prefillRawAudioInput(
       byte[] audio, int batch_size, int n_channels, int n_samples);
 
   /**
-   * Prefill a multimodal Module with the given text input.
+   * Prefill the KV cache with the given text prompt.
    *
    * @param prompt The text prompt to prefill.
-   * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer
-   *     exposed to user.
+   * @return 0 on success
    * @throws RuntimeException if the prefill failed
    */
   @Experimental
   public long prefillPrompt(String prompt) {
-    int nativeResult = appendTextInput(prompt);
+    int nativeResult = prefillTextInput(prompt);
     if (nativeResult != 0) {
       throw new RuntimeException("Prefill failed with error code: " + nativeResult);
     }
     return 0;
   }
 
   // returns status
-  private native int appendTextInput(String prompt);
+  private native int prefillTextInput(String prompt);
 
   /**
    * Reset the context of the LLM. This will clear the KV cache and reset the state of the LLM.