1 files changed, 69 insertions, 248 deletions

diff --git a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc
index 31f7b58fbc..fdcf00a0a0 100644
--- a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc
+++ b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc
@@ -16,9 +16,6 @@ limitations under the License.
 #include "tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h"
 namespace {
 
-const int kByteBufferValue = 999;
-const int kBufferSize = 256;
-
 tflite::Interpreter* convertLongToInterpreter(JNIEnv* env, jlong handle) {
   if (handle == 0) {
     throwException(env, kIllegalArgumentException,
@@ -62,22 +59,6 @@ std::vector<int> convertJIntArrayToVector(JNIEnv* env, jintArray inputs) {
   return outputs;
 }
 
-bool isByteBuffer(jint data_type) { return data_type == kByteBufferValue; }
-
-TfLiteType resolveDataType(jint data_type) {
-  switch (data_type) {
-    case 1:
-      return kTfLiteFloat32;
-    case 2:
-      return kTfLiteInt32;
-    case 3:
-      return kTfLiteUInt8;
-    case 4:
-      return kTfLiteInt64;
-    default:
-      return kTfLiteNoType;
-  }
-}
 
 int getDataType(TfLiteType data_type) {
   switch (data_type) {
@@ -108,64 +89,6 @@ void printDims(char* buffer, int max_size, int* dims, int num_dims) {
   }
 }
 
-TfLiteStatus checkInputs(JNIEnv* env, tflite::Interpreter* interpreter,
-                         const int input_size, jintArray data_types,
-                         jintArray nums_of_bytes, jobjectArray values,
-                         jobjectArray sizes) {
-  if (input_size != interpreter->inputs().size()) {
-    throwException(env, kIllegalArgumentException,
-                   "Input error: Expected num of inputs is %d but got %d",
-                   interpreter->inputs().size(), input_size);
-    return kTfLiteError;
-  }
-  if (input_size != env->GetArrayLength(data_types) ||
-      input_size != env->GetArrayLength(nums_of_bytes) ||
-      input_size != env->GetArrayLength(values)) {
-    throwException(env, kIllegalArgumentException,
-                   "Internal error: Arrays in arguments should be of the same "
-                   "length, but got %d sizes, %d data_types, %d nums_of_bytes, "
-                   "and %d values",
-                   input_size, env->GetArrayLength(data_types),
-                   env->GetArrayLength(nums_of_bytes),
-                   env->GetArrayLength(values));
-    return kTfLiteError;
-  }
-  for (int i = 0; i < input_size; ++i) {
-    int input_idx = interpreter->inputs()[i];
-    TfLiteTensor* target = interpreter->tensor(input_idx);
-    jintArray dims =
-        static_cast<jintArray>(env->GetObjectArrayElement(sizes, i));
-    int num_dims = static_cast<int>(env->GetArrayLength(dims));
-    if (target->dims->size != num_dims) {
-      throwException(env, kIllegalArgumentException,
-                     "Input error: %d-th input should have %d dimensions, but "
-                     "found %d dimensions",
-                     i, target->dims->size, num_dims);
-      return kTfLiteError;
-    }
-    jint* ptr = env->GetIntArrayElements(dims, nullptr);
-    for (int j = 1; j < num_dims; ++j) {
-      if (target->dims->data[j] != ptr[j]) {
-        std::unique_ptr<char[]> expected_dims(new char[kBufferSize]);
-        std::unique_ptr<char[]> obtained_dims(new char[kBufferSize]);
-        printDims(expected_dims.get(), kBufferSize, target->dims->data,
-                  num_dims);
-        printDims(obtained_dims.get(), kBufferSize, ptr, num_dims);
-        throwException(env, kIllegalArgumentException,
-                       "Input error: %d-th input dimension should be [%s], but "
-                       "found [%s]",
-                       i, expected_dims.get(), obtained_dims.get());
-        env->ReleaseIntArrayElements(dims, ptr, JNI_ABORT);
-        return kTfLiteError;
-      }
-    }
-    env->ReleaseIntArrayElements(dims, ptr, JNI_ABORT);
-    env->DeleteLocalRef(dims);
-    if (env->ExceptionCheck()) return kTfLiteError;
-  }
-  return kTfLiteOk;
-}
-
 // Checks whether there is any difference between dimensions of a tensor and a
 // given dimensions. Returns true if there is difference, else false.
 bool areDimsDifferent(JNIEnv* env, TfLiteTensor* tensor, jintArray dims) {
@@ -188,74 +111,6 @@ bool areDimsDifferent(JNIEnv* env, TfLiteTensor* tensor, jintArray dims) {
   return false;
 }
 
-bool areInputDimensionsTheSame(JNIEnv* env, tflite::Interpreter* interpreter,
-                               int input_size, jobjectArray sizes) {
-  if (interpreter->inputs().size() != input_size) {
-    return false;
-  }
-  for (int i = 0; i < input_size; ++i) {
-    int input_idx = interpreter->inputs()[i];
-    jintArray dims =
-        static_cast<jintArray>(env->GetObjectArrayElement(sizes, i));
-    TfLiteTensor* target = interpreter->tensor(input_idx);
-    if (areDimsDifferent(env, target, dims)) return false;
-    env->DeleteLocalRef(dims);
-    if (env->ExceptionCheck()) return false;
-  }
-  return true;
-}
-
-TfLiteStatus resizeInputs(JNIEnv* env, tflite::Interpreter* interpreter,
-                          int input_size, jobjectArray sizes) {
-  for (int i = 0; i < input_size; ++i) {
-    int input_idx = interpreter->inputs()[i];
-    jintArray dims =
-        static_cast<jintArray>(env->GetObjectArrayElement(sizes, i));
-    TfLiteStatus status = interpreter->ResizeInputTensor(
-        input_idx, convertJIntArrayToVector(env, dims));
-    if (status != kTfLiteOk) {
-      return status;
-    }
-    env->DeleteLocalRef(dims);
-    if (env->ExceptionCheck()) return kTfLiteError;
-  }
-  return kTfLiteOk;
-}
-
-TfLiteStatus setInputs(JNIEnv* env, tflite::Interpreter* interpreter,
-                       int input_size, jintArray data_types,
-                       jintArray nums_of_bytes, jobjectArray values) {
-  jint* data_type = env->GetIntArrayElements(data_types, nullptr);
-  jint* num_bytes = env->GetIntArrayElements(nums_of_bytes, nullptr);
-  for (int i = 0; i < input_size; ++i) {
-    int input_idx = interpreter->inputs()[i];
-    TfLiteTensor* target = interpreter->tensor(input_idx);
-    jobject value = env->GetObjectArrayElement(values, i);
-    bool is_byte_buffer = isByteBuffer(data_type[i]);
-    if (is_byte_buffer) {
-      writeByteBuffer(env, value, &(target->data.raw),
-                      static_cast<int>(num_bytes[i]));
-    } else {
-      TfLiteType type = resolveDataType(data_type[i]);
-      if (type != target->type) {
-        throwException(env, kIllegalArgumentException,
-                       "Input error: DataType (%d) of input data does not "
-                       "match with the DataType (%d) of model inputs.",
-                       type, target->type);
-        return kTfLiteError;
-      }
-      writeMultiDimensionalArray(env, value, target->type, target->dims->size,
-                                 &(target->data.raw),
-                                 static_cast<int>(num_bytes[i]));
-    }
-    env->DeleteLocalRef(value);
-    if (env->ExceptionCheck()) return kTfLiteError;
-  }
-  env->ReleaseIntArrayElements(data_types, data_type, JNI_ABORT);
-  env->ReleaseIntArrayElements(nums_of_bytes, num_bytes, JNI_ABORT);
-  return kTfLiteOk;
-}
-
 // TODO(yichengfan): evaluate the benefit to use tflite verifier.
 bool VerifyModel(const void* buf, size_t len) {
   flatbuffers::Verifier verifier(static_cast<const uint8_t*>(buf), len);
@@ -287,6 +142,64 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputNames(JNIEnv* env,
   return names;
 }
 
+JNIEXPORT void JNICALL
+Java_org_tensorflow_lite_NativeInterpreterWrapper_allocateTensors(
+    JNIEnv* env, jclass clazz, jlong handle, jlong error_handle) {
+  tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle);
+  if (interpreter == nullptr) return;
+  BufferErrorReporter* error_reporter =
+      convertLongToErrorReporter(env, error_handle);
+  if (error_reporter == nullptr) return;
+
+  if (interpreter->AllocateTensors() != kTfLiteOk) {
+    throwException(
+        env, kIllegalStateException,
+        "Internal error: Unexpected failure when preparing tensor allocations:"
+        " %s",
+        error_reporter->CachedErrorMessage());
+  }
+}
+
+JNIEXPORT jlong JNICALL
+Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputTensor(JNIEnv* env,
+                                                                 jclass clazz,
+                                                                 jlong handle,
+                                                                 jint index) {
+  tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle);
+  if (interpreter == nullptr) return 0;
+  return reinterpret_cast<jlong>(
+      interpreter->tensor(interpreter->inputs()[index]));
+}
+
+JNIEXPORT jlong JNICALL
+Java_org_tensorflow_lite_NativeInterpreterWrapper_getOutputTensor(JNIEnv* env,
+                                                                  jclass clazz,
+                                                                  jlong handle,
+                                                                  jint index) {
+  tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle);
+  if (interpreter == nullptr) return 0;
+  return reinterpret_cast<jlong>(
+      interpreter->tensor(interpreter->outputs()[index]));
+}
+
+JNIEXPORT jint JNICALL
+Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputCount(JNIEnv* env,
+                                                                jclass clazz,
+                                                                jlong handle) {
+  tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle);
+  if (interpreter == nullptr) return 0;
+  return static_cast<jint>(interpreter->inputs().size());
+}
+
+JNIEXPORT jint JNICALL
+Java_org_tensorflow_lite_NativeInterpreterWrapper_getOutputCount(JNIEnv* env,
+                                                                 jclass clazz,
+                                                                 jlong handle) {
+  tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle);
+  if (interpreter == nullptr) return 0;
+  return static_cast<jint>(interpreter->outputs().size());
+}
+
 JNIEXPORT jobjectArray JNICALL
 Java_org_tensorflow_lite_NativeInterpreterWrapper_getOutputNames(JNIEnv* env,
                                                                  jclass clazz,
@@ -424,124 +337,32 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createInterpreter(
   // allocates memory
   status = interpreter->AllocateTensors();
   if (status != kTfLiteOk) {
-    throwException(env, kNullPointerException,
-                   "Internal error: Cannot allocate memory for the interpreter:"
-                   " %s",
-                   error_reporter->CachedErrorMessage());
+    throwException(
+        env, kIllegalStateException,
+        "Internal error: Unexpected failure when preparing tensor allocations:"
+        " %s",
+        error_reporter->CachedErrorMessage());
     return 0;
   }
   return reinterpret_cast<jlong>(interpreter.release());
 }
 
 // Sets inputs, runs inference, and returns outputs as long handles.
-JNIEXPORT jlongArray JNICALL
-Java_org_tensorflow_lite_NativeInterpreterWrapper_run(
-    JNIEnv* env, jclass clazz, jlong interpreter_handle, jlong error_handle,
-    jobjectArray sizes, jintArray data_types, jintArray nums_of_bytes,
-    jobjectArray values, jobject wrapper, jboolean memory_allocated) {
+JNIEXPORT void JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_run(
+    JNIEnv* env, jclass clazz, jlong interpreter_handle, jlong error_handle) {
   tflite::Interpreter* interpreter =
       convertLongToInterpreter(env, interpreter_handle);
-  if (interpreter == nullptr) return nullptr;
+  if (interpreter == nullptr) return;
   BufferErrorReporter* error_reporter =
       convertLongToErrorReporter(env, error_handle);
-  if (error_reporter == nullptr) return nullptr;
-  const int input_size = env->GetArrayLength(sizes);
-  // validates inputs
-  TfLiteStatus status = checkInputs(env, interpreter, input_size, data_types,
-                                    nums_of_bytes, values, sizes);
-  if (status != kTfLiteOk) return nullptr;
-  if (!memory_allocated ||
-      !areInputDimensionsTheSame(env, interpreter, input_size, sizes)) {
-    // resizes inputs
-    status = resizeInputs(env, interpreter, input_size, sizes);
-    if (status != kTfLiteOk) {
-      throwException(env, kNullPointerException,
-                     "Internal error: Can not resize the input: %s",
-                     error_reporter->CachedErrorMessage());
-      return nullptr;
-    }
-    // allocates memory
-    status = interpreter->AllocateTensors();
-    if (status != kTfLiteOk) {
-      throwException(env, kNullPointerException,
-                     "Internal error: Can not allocate memory for the given "
-                     "inputs: %s",
-                     error_reporter->CachedErrorMessage());
-      return nullptr;
-    }
-  }
-  // sets inputs
-  status = setInputs(env, interpreter, input_size, data_types, nums_of_bytes,
-                     values);
-  if (status != kTfLiteOk) return nullptr;
-  timespec beforeInference = ::tflite::getCurrentTime();
-  // runs inference
+  if (error_reporter == nullptr) return;
+
   if (interpreter->Invoke() != kTfLiteOk) {
     throwException(env, kIllegalArgumentException,
                    "Internal error: Failed to run on the given Interpreter: %s",
                    error_reporter->CachedErrorMessage());
-    return nullptr;
-  }
-  timespec afterInference = ::tflite::getCurrentTime();
-  jclass wrapper_clazz = env->GetObjectClass(wrapper);
-  jfieldID fid =
-      env->GetFieldID(wrapper_clazz, "inferenceDurationNanoseconds", "J");
-  if (env->ExceptionCheck()) {
-    env->ExceptionClear();
-  } else if (fid != nullptr) {
-    env->SetLongField(
-        wrapper, fid,
-        ::tflite::timespec_diff_nanoseconds(&beforeInference, &afterInference));
-  }
-  // returns outputs
-  const std::vector<int>& results = interpreter->outputs();
-  if (results.empty()) {
-    throwException(
-        env, kIllegalArgumentException,
-        "Internal error: The Interpreter does not have any outputs.");
-    return nullptr;
-  }
-  jlongArray outputs = env->NewLongArray(results.size());
-  size_t size = results.size();
-  for (int i = 0; i < size; ++i) {
-    TfLiteTensor* source = interpreter->tensor(results[i]);
-    jlong output = reinterpret_cast<jlong>(source);
-    env->SetLongArrayRegion(outputs, i, 1, &output);
-  }
-  return outputs;
-}
-
-JNIEXPORT jintArray JNICALL
-Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputDims(
-    JNIEnv* env, jclass clazz, jlong handle, jint input_idx, jint num_bytes) {
-  tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle);
-  if (interpreter == nullptr) return nullptr;
-  const int idx = static_cast<int>(input_idx);
-  if (input_idx < 0 || input_idx >= interpreter->inputs().size()) {
-    throwException(env, kIllegalArgumentException,
-                   "Input error: Out of range: Failed to get %d-th input out of"
-                   " %d inputs",
-                   input_idx, interpreter->inputs().size());
-    return nullptr;
-  }
-  TfLiteTensor* target = interpreter->tensor(interpreter->inputs()[idx]);
-  int size = target->dims->size;
-  if (num_bytes >= 0) {  // verifies num of bytes matches if num_bytes if valid.
-    int expected_num_bytes = elementByteSize(target->type);
-    for (int i = 0; i < size; ++i) {
-      expected_num_bytes *= target->dims->data[i];
-    }
-    if (num_bytes != expected_num_bytes) {
-      throwException(env, kIllegalArgumentException,
-                     "Input error: Failed to get input dimensions. %d-th input "
-                     "should have %d bytes, but found %d bytes.",
-                     idx, expected_num_bytes, num_bytes);
-      return nullptr;
-    }
+    return;
   }
-  jintArray outputs = env->NewIntArray(size);
-  env->SetIntArrayRegion(outputs, 0, size, &(target->dims->data[0]));
-  return outputs;
 }
 
 JNIEXPORT jint JNICALL