Use InterpreterBuilder::AddDelegate (attempt #2).

Use the non-experimental method InterpreterBuilder::AddDelegate to
add delegates before constructing the interpreter, rather
than using the experimental method Interpreter::ModifyGraphWithDelegate
to set delegates after the interpreter is already constructed.

A significant drawback of this refactoring is that if any delegates
are enabled, we end up having to construct the interpreter twice:
once without delegates, in order to determine whether the model
contains any unresolved flex ops in order to decide whether to
enable the flex delegate, and then again with all the delegates.

PiperOrigin-RevId: 401003838
Change-Id: I8fdf8a4728a7a98677c5f8cea895447df234310b

Author: fergushenderson

tensorflow/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java

@@ -72,25 +72,30 @@ private void init(long errorHandle, long modelHandle, InterpreterImpl.Options op
     }
     this.errorHandle = errorHandle;
     this.modelHandle = modelHandle;
-    // First create the interpreter without delegates.
-    // We need this in order to figure out whether the model contains any unresolved flex ops.
+    // First create the interpreter without delegates.  We need an interpreter in order to figure
+    // out whether the model contains any unresolved flex ops, and creating the interpreter with
+    // delegates might fail if there are any unresolved flex ops.
     // (Alternatively, we could determine this without needing to recreate the interpreter
     // by passing the tflite::Model in to here, and then traversing that?)
-    this.interpreterHandle = createInterpreter(modelHandle, errorHandle, options.numThreads);
+    ArrayList<Long> delegateHandles = new ArrayList<Long>();
+    this.interpreterHandle =
+        createInterpreter(modelHandle, errorHandle, options.numThreads, delegateHandles);
     this.originalGraphHasUnresolvedFlexOp = hasUnresolvedFlexOp(interpreterHandle);
     addDelegates(options);
-    // TODO(b/187920750): uncomment this when createInterpreter is modified to use
-    // InterpreterBuilder::AddDelegate.
-    // if (!delegates.isEmpty()) {
-    //   // If there are any delegates enabled, recreate the interpreter with those delegates.
-    //   delete(/* errorHandle= */ 0, /* modelHandle= */ 0, this.interpreterHandle);
-    //   this.interpreterHandle = createInterpreter(modelHandle, errorHandle, options.numThreads);
-    // }
+    delegateHandles.ensureCapacity(delegates.size());
+    for (Delegate delegate : delegates) {
+      delegateHandles.add(Long.valueOf(delegate.getNativeHandle()));
+    }
+    if (!delegateHandles.isEmpty()) {
+      // If there are any delegates enabled, recreate the interpreter with those delegates.
+      delete(/* errorHandle= */ 0, /* modelHandle= */ 0, this.interpreterHandle);
+      this.interpreterHandle =
+          createInterpreter(modelHandle, errorHandle, options.numThreads, delegateHandles);
+    }
     if (options.allowFp16PrecisionForFp32 != null) {
       allowFp16PrecisionForFp32(
           interpreterHandle, options.allowFp16PrecisionForFp32.booleanValue());
     }
-    applyDelegates(options);
     if (options.allowBufferHandleOutput != null) {
       allowBufferHandleOutput(interpreterHandle, options.allowBufferHandleOutput.booleanValue());
     }
@@ -524,13 +529,6 @@ private void maybeAddXnnpackDelegate(InterpreterImpl.Options options) {
     }
   }
 
-  // Apply all the delegates specified in this.delegates.
-  private void applyDelegates(InterpreterImpl.Options options) {
-    for (Delegate delegate : delegates) {
-      applyDelegate(interpreterHandle, errorHandle, delegate.getNativeHandle());
-    }
-  }
-
   private NativeSignatureRunnerWrapper getSignatureRunnerWrapper(String signatureKey) {
     if (signatureRunnerMap == null) {
       signatureRunnerMap = new HashMap<>();
@@ -643,10 +641,8 @@ private static native XnnpackDelegate createXNNPACKDelegate(
 
   private static native long createModelWithBuffer(ByteBuffer modelBuffer, long errorHandle);
 
-  private static native long createInterpreter(long modelHandle, long errorHandle, int numThreads);
-
-  private static native void applyDelegate(
-      long interpreterHandle, long errorHandle, long delegateHandle);
+  private static native long createInterpreter(
+      long modelHandle, long errorHandle, int numThreads, List<Long> delegateHandles);
 
   private static native void resetVariableTensors(long interpreterHandle, long errorHandle);
 

tensorflow/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc

@@ -494,25 +494,121 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createModelWithBuffer(
 JNIEXPORT jlong JNICALL
 Java_org_tensorflow_lite_NativeInterpreterWrapper_createInterpreter(
     JNIEnv* env, jclass clazz, jlong model_handle, jlong error_handle,
-    jint num_threads) {
+    jint num_threads, jobject delegate_handle_list) {
   if (!tflite::jni::CheckJniInitializedOrThrow(env)) return 0;
 
+  static jclass list_class = env->FindClass("java/util/List");
+  if (list_class == nullptr) {
+    if (!env->ExceptionCheck()) {
+      ThrowException(env, tflite::jni::kUnsupportedOperationException,
+                     "Internal error: Can't find java.util.List class.");
+    }
+    return 0;
+  }
+  static jmethodID list_size_method =
+      env->GetMethodID(list_class, "size", "()I");
+  if (list_size_method == nullptr) {
+    if (!env->ExceptionCheck()) {
+      ThrowException(env, tflite::jni::kUnsupportedOperationException,
+                     "Internal error: Can't find java.util.List.size method.");
+    }
+    return 0;
+  }
+  static jmethodID list_get_method =
+      env->GetMethodID(list_class, "get", "(I)Ljava/lang/Object;");
+  if (list_get_method == nullptr) {
+    if (!env->ExceptionCheck()) {
+      ThrowException(env, tflite::jni::kUnsupportedOperationException,
+                     "Internal error: Can't find java.util.List.get method.");
+    }
+    return 0;
+  }
+  static jclass long_class = env->FindClass("java/lang/Long");
+  if (long_class == nullptr) {
+    if (!env->ExceptionCheck()) {
+      ThrowException(env, tflite::jni::kUnsupportedOperationException,
+                     "Internal error: "
+                     "Can't find java.lang.Long class.");
+    }
+    return 0;
+  }
+  static jmethodID long_value_method =
+      env->GetMethodID(long_class, "longValue", "()J");
+  if (long_value_method == nullptr) {
+    if (!env->ExceptionCheck()) {
+      ThrowException(env, tflite::jni::kUnsupportedOperationException,
+                     "Internal error: "
+                     "Can't find java.lang.Long longValue method.");
+    }
+    return 0;
+  }
+
   FlatBufferModel* model = convertLongToModel(env, model_handle);
   if (model == nullptr) return 0;
+
   BufferErrorReporter* error_reporter =
       convertLongToErrorReporter(env, error_handle);
   if (error_reporter == nullptr) return 0;
+
   std::unique_ptr<OpResolver> resolver = tflite_shims::CreateOpResolver();
+
   InterpreterBuilder interpreter_builder(*model, *resolver);
   interpreter_builder.SetNumThreads(static_cast<int>(num_threads));
+
+  // Add delegate_list to interpreter_builder.
+
+  // Java: int size = delegate_list.size();
+  jint size = env->CallIntMethod(delegate_handle_list, list_size_method);
+  for (jint i = 0; i < size; ++i) {
+    // Java: Long jdelegate_handle = delegate_handle_list->get(i);
+    jobject jdelegate_handle =
+        env->CallObjectMethod(delegate_handle_list, list_get_method, i);
+    if (jdelegate_handle == nullptr) {
+      ThrowException(env, tflite::jni::kIllegalArgumentException,
+                     "Internal error: null Delegate handle");
+      return 0;
+    }
+    // Java: long delegate_handle = jdelegate_handle.longValue();
+    jlong delegate_handle =
+        env->CallLongMethod(jdelegate_handle, long_value_method);
+    if (delegate_handle == 0) {
+      ThrowException(env, tflite::jni::kIllegalArgumentException,
+                     "Internal error: Found invalid handle");
+      return 0;
+    }
+    auto delegate = reinterpret_cast<TfLiteOpaqueDelegate*>(delegate_handle);
+    interpreter_builder.AddDelegate(delegate);
+  }
+
+  // Create the Interpreter.
   std::unique_ptr<Interpreter> interpreter;
   TfLiteStatus status = interpreter_builder(&interpreter);
   if (status != kTfLiteOk) {
-    ThrowException(env, tflite::jni::kIllegalArgumentException,
-                   "Internal error: Cannot create interpreter: %s",
-                   error_reporter->CachedErrorMessage());
+    if (status == kTfLiteDelegateError) {
+      ThrowException(env, tflite::jni::kIllegalArgumentException,
+                     "Internal error: Failed to apply delegate: %s",
+                     error_reporter->CachedErrorMessage());
+    } else if (status == kTfLiteApplicationError) {
+      ThrowException(env, tflite::jni::kIllegalArgumentException,
+                     "Internal error: Error applying delegate: %s",
+                     error_reporter->CachedErrorMessage());
+    } else {
+      const char* error_message = error_reporter->CachedErrorMessage();
+      if (std::strcmp(
+              error_message,
+              "Restored original execution plan after delegate application "
+              "failure.") == 0) {
+        ThrowException(env, tflite::jni::kIllegalArgumentException,
+                       "Internal error: Failed to apply delegate.");
+      } else {
+        ThrowException(env, tflite::jni::kIllegalArgumentException,
+                       "Internal error: Cannot create interpreter: %s",
+                       error_message);
+      }
+    }
     return 0;
   }
+
   // Note that tensor allocation is performed explicitly by the owning Java
   // NativeInterpreterWrapper instance.
   return reinterpret_cast<jlong>(interpreter.release());
@@ -598,30 +694,6 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_resizeInput(
   return is_changed ? JNI_TRUE : JNI_FALSE;
 }
 
-JNIEXPORT void JNICALL
-Java_org_tensorflow_lite_NativeInterpreterWrapper_applyDelegate(
-    JNIEnv* env, jclass clazz, jlong interpreter_handle, jlong error_handle,
-    jlong delegate_handle) {
-  if (!tflite::jni::CheckJniInitializedOrThrow(env)) return;
-
-  Interpreter* interpreter = convertLongToInterpreter(env, interpreter_handle);
-  if (interpreter == nullptr) return;
-
-  BufferErrorReporter* error_reporter =
-      convertLongToErrorReporter(env, error_handle);
-  if (error_reporter == nullptr) return;
-
-  TfLiteOpaqueDelegate* delegate = convertLongToDelegate(env, delegate_handle);
-  if (delegate == nullptr) return;
-
-  TfLiteStatus status = interpreter->ModifyGraphWithDelegate(delegate);
-  if (status != kTfLiteOk) {
-    ThrowException(env, tflite::jni::kIllegalArgumentException,
-                   "Internal error: Failed to apply delegate: %s",
-                   error_reporter->CachedErrorMessage());
-  }
-}
-
 JNIEXPORT jlong JNICALL
 Java_org_tensorflow_lite_NativeInterpreterWrapper_createCancellationFlag(
     JNIEnv* env, jclass clazz, jlong interpreter_handle) {