[Java API] Infer and updates to NeuropodTensor (uber#421)

### Summary:
This is a next step in implementation of Java API that includes Inference method and necessary updates to other modules.
Also several JNI C++ methods changed to use std::string instead of "const char*".
  
### Test Plan:
Unit test covers Inference of TF addition model.

source/neuropod/bindings/java/src/main/java/com/uber/neuropod/Neuropod.java

@@ -75,7 +75,7 @@ public String getPlatform() {
      * @return the inference result
      */
     public Map<String, NeuropodTensor> infer(Map<String, NeuropodTensor> inputs) {
-        return null;
+        return infer(inputs, null);
     }
 
     /**
@@ -87,7 +87,7 @@ public Map<String, NeuropodTensor> infer(Map<String, NeuropodTensor> inputs) {
      * @return the inference result
      */
     public Map<String, NeuropodTensor> infer(Map<String, NeuropodTensor> inputs, List<String> requestedOutputs) {
-        return null;
+        return nativeInfer(inputs.entrySet().toArray(), requestedOutputs, super.getNativeHandle());
     }
 
     /**
@@ -149,6 +149,9 @@ public static NeuropodTensorAllocator getGenericTensorAllocator() {
 
     private static native long nativeGetGenericAllocator();
 
+    private static native Map<String, NeuropodTensor> nativeInfer(Object[] inputs, List<String> requestedOutputs,
+                                                                  long modelHandle);
+
     @Override
     protected native void nativeDelete(long handle);
 }

source/neuropod/bindings/java/src/main/java/com/uber/neuropod/NeuropodTensor.java

@@ -29,13 +29,23 @@ public class NeuropodTensor extends NativeClass implements Serializable {
 
     protected ByteBuffer buffer;
 
+    // This flag is used to separate NeuropodTensor that is created as Input (from java)
+    // or Output (from JNI as result of inference). This is necessary because tensor buffer
+    // is allocated differently and it is not safe to expose data directly to user.
     private final boolean isFromJava;
 
     // Constructor for NeuropodTensor created in Java side.
     protected NeuropodTensor() {
         isFromJava = true;
     }
 
+    // Constructor for NeuropodTensor created in C++ side.
+    protected NeuropodTensor(long handle) {
+        super(handle);
+        isFromJava = false;
+        buffer = nativeGetBuffer(handle).order(ByteOrder.nativeOrder());
+    }
+
     /**
      * Get the dims array which represnents the shape of a tensor.
      *
@@ -66,9 +76,9 @@ public TensorType getTensorType() {
     /**
      * Flatten the tensor data and convert it to a long buffer.
      * <p>
-     * Can only be used when the tensor is INT64_TENSOR. Will trigger
-     * a copy if the tensor is created by infer method. Otherwise will
-     * not trigger a copy.
+     * Can only be used when the tensor is INT64_TENSOR.
+     *
+     * Note: It returns a buffer that is valid even after Tensor is closed.
      *
      * @return the IntBuffer
      */
@@ -77,7 +87,28 @@ public LongBuffer toLongBuffer() {
         if (isFromJava) {
             return buffer.asLongBuffer();
         }
-        return null;
+        LongBuffer ret = LongBuffer.allocate((int)getNumberOfElements()).put(buffer.asLongBuffer());
+        ret.rewind();
+        return ret;
+    }
+
+    /**
+     * Flatten the tensor data and convert it to a float buffer.
+     * <p>
+     * Can only be used when the tensor is FLOAT_TENSOR.
+     *
+     * Note: It returns a buffer that is valid even after Tensor is closed.
+     *
+     * @return the FloatBuffer
+     */
+    public FloatBuffer toFloatBuffer() {
+        checkType(TensorType.FLOAT_TENSOR);
+        if (isFromJava) {
+            return buffer.asFloatBuffer();
+        }
+        FloatBuffer ret = FloatBuffer.allocate((int)getNumberOfElements()).put(buffer.asFloatBuffer());
+        ret.rewind();
+        return ret;
     }
 
      /**
@@ -90,21 +121,12 @@ public LongBuffer toLongBuffer() {
       */
      public IntBuffer toIntBuffer() {return null;}
 
-     /**
-      * Flatten the tensor data and convert it to a float buffer.
-      * <p>
-      * Can only be used when the tensor is FLOAT_TENSOR. Will trigger
-      * a copy.
-      *
-      * @return the FloatBuffer
-      */
-     public FloatBuffer toFloatBuffer() {return null;}
-
      /**
       * Flatten the tensor data and convert it to a double buffer.
       * <p>
-      * Can only be used when the tensor is DOUBLE_TENSOR. Will trigger
-      * a copy.
+      * Can only be used when the tensor is DOUBLE_TENSOR.
+      *
+      * Note: It returns a buffer that is valid even after Tensor is closed.
       *
       * @return the DoubleBuffer
       */
@@ -188,6 +210,11 @@ private void checkType(TensorType type) {
         }
     }
 
+    // Easier for the JNI side to call methods of super class.
+    private long getHandle() {
+        return super.getNativeHandle();
+    }
+
     @Override
     protected void nativeDelete(long handle) throws NeuropodJNIException {
         nativeDoDelete(handle);
@@ -201,4 +228,6 @@ protected void nativeDelete(long handle) throws NeuropodJNIException {
     private static native TensorType nativeGetTensorType(long nativeHandle) throws NeuropodJNIException;
 
     private static native long nativeGetNumberOfElements(long nativeHandle) throws NeuropodJNIException;
+
+    private static native ByteBuffer nativeGetBuffer(long nativeHandle);
 }

source/neuropod/bindings/java/src/main/native/com_uber_neuropod_Neuropod.cc

@@ -26,6 +26,7 @@ limitations under the License.
 #include <memory>
 #include <sstream>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include <jni.h>
@@ -215,3 +216,64 @@ JNIEXPORT jlong JNICALL Java_com_uber_neuropod_Neuropod_nativeGetGenericAllocato
     }
     return reinterpret_cast<jlong>(nullptr);
 }
+
+JNIEXPORT jobject JNICALL Java_com_uber_neuropod_Neuropod_nativeInfer(
+    JNIEnv *env, jclass, jobjectArray entryArray, jobject requestedOutputsJava, jlong modelHandle)
+{
+    try
+    {
+        // Prepare requestedOutputs
+        std::vector<std::string> requestedOutputs;
+        if (requestedOutputsJava != nullptr)
+        {
+            jsize size = env->CallIntMethod(requestedOutputsJava, java_util_ArrayList_size);
+            for (jsize i = 0; i < size; i++)
+            {
+                jstring element =
+                    static_cast<jstring>(env->CallObjectMethod(requestedOutputsJava, java_util_ArrayList_get, i));
+                requestedOutputs.emplace_back(toString(env, element));
+                env->DeleteLocalRef(element);
+            }
+        }
+
+        // Fill in NeuropodValueMap
+        jsize                      entrySize = env->GetArrayLength(entryArray);
+        neuropod::NeuropodValueMap nativeMap;
+        for (jsize i = 0; i < entrySize; i++)
+        {
+            jobject     entry = env->GetObjectArrayElement(entryArray, i);
+            std::string key =
+                toString(env, static_cast<jstring>(env->CallObjectMethod(entry, java_util_Map_Entry_getKey)));
+            jobject value        = env->CallObjectMethod(entry, java_util_Map_Entry_getValue);
+            jlong   tensorHandle = env->CallLongMethod(value, com_uber_neuropod_NeuropodTensor_getHandle);
+            if (tensorHandle == 0)
+            {
+                throw std::runtime_error("unexpected NULL tensor handle");
+            }
+            nativeMap.insert(
+                std::make_pair(key, *reinterpret_cast<std::shared_ptr<neuropod::NeuropodValue> *>(tensorHandle)));
+            env->DeleteLocalRef(entry);
+            env->DeleteLocalRef(value);
+        }
+
+        auto model       = reinterpret_cast<neuropod::Neuropod *>(modelHandle);
+        auto inferredMap = model->infer(nativeMap, requestedOutputs);
+
+        // Put data to Java Map
+        auto ret = env->NewObject(java_util_HashMap, java_util_HashMap_);
+        for (auto &entry : *inferredMap)
+        {
+            jobject javaTensor = env->NewObject(com_uber_neuropod_NeuropodTensor,
+                                                com_uber_neuropod_NeuropodTensor_,
+                                                reinterpret_cast<jlong>(toHeap(entry.second)));
+            env->CallObjectMethod(ret, java_util_HashMap_put, env->NewStringUTF(entry.first.c_str()), javaTensor);
+            env->DeleteLocalRef(javaTensor);
+        }
+        return ret;
+    }
+    catch (const std::exception &e)
+    {
+        throwJavaException(env, e.what());
+    }
+    return nullptr;
+}

source/neuropod/bindings/java/src/main/native/com_uber_neuropod_Neuropod.h

@@ -81,6 +81,13 @@ JNIEXPORT jlong JNICALL Java_com_uber_neuropod_Neuropod_nativeGetAllocator(JNIEn
  */
 JNIEXPORT jlong JNICALL Java_com_uber_neuropod_Neuropod_nativeGetGenericAllocator(JNIEnv *, jclass);
 
+/*
+ * Class:     com_uber_neuropod_Neuropod
+ * Method:    nativeInfer
+ * Signature: ([Ljava/lang/Object;Ljava/util/List;J)Ljava/util/Map;
+ */
+JNIEXPORT jobject JNICALL Java_com_uber_neuropod_Neuropod_nativeInfer(JNIEnv *, jclass, jobjectArray, jobject, jlong);
+
 /*
  * Class:     com_uber_neuropod_Neuropod
  * Method:    nativeDelete

source/neuropod/bindings/java/src/main/native/com_uber_neuropod_NeuropodTensor.cc

@@ -26,6 +26,38 @@ JNIEXPORT void JNICALL Java_com_uber_neuropod_NeuropodTensor_nativeDoDelete(JNIE
     }
 }
 
+JNIEXPORT jobject JNICALL Java_com_uber_neuropod_NeuropodTensor_nativeGetBuffer(JNIEnv *env, jclass, jlong nativeHandle)
+{
+    try
+    {
+        auto neuropodTensor =
+            (*reinterpret_cast<std::shared_ptr<neuropod::NeuropodValue> *>(nativeHandle))->as_tensor();
+        auto tensorType = neuropodTensor->get_tensor_type();
+        switch (tensorType)
+        {
+        case neuropod::FLOAT_TENSOR: {
+            return createDirectBuffer<float>(env, neuropodTensor);
+        }
+        case neuropod::DOUBLE_TENSOR: {
+            return createDirectBuffer<double>(env, neuropodTensor);
+        }
+        case neuropod::INT32_TENSOR: {
+            return createDirectBuffer<int32_t>(env, neuropodTensor);
+        }
+        case neuropod::INT64_TENSOR: {
+            return createDirectBuffer<int64_t>(env, neuropodTensor);
+        }
+        default:
+            throw std::runtime_error("Unsupported tensor type: " + tensorTypeToString(tensorType));
+        }
+    }
+    catch (const std::exception &e)
+    {
+        throwJavaException(env, e.what());
+    }
+    return nullptr;
+}
+
 JNIEXPORT jlongArray JNICALL Java_com_uber_neuropod_NeuropodTensor_nativeGetDims(JNIEnv *env, jclass, jlong handle)
 {
     try

source/neuropod/bindings/java/src/main/native/com_uber_neuropod_NeuropodTensor.h

@@ -29,6 +29,13 @@ extern "C" {
  */
 JNIEXPORT void JNICALL Java_com_uber_neuropod_NeuropodTensor_nativeDoDelete(JNIEnv *, jobject, jlong);
 
+/*
+ * Class:     com_uber_neuropod_NeuropodTensor
+ * Method:    nativeGetBuffer
+ * Signature: (J)Ljava/nio/ByteBuffer;
+ */
+JNIEXPORT jobject JNICALL Java_com_uber_neuropod_NeuropodTensor_nativeGetBuffer(JNIEnv *, jclass, jlong);
+
 /*
  * Class:     com_uber_neuropod_NeuropodTensor
  * Method:    nativeGetDims

source/neuropod/bindings/java/src/main/native/jclass_register.cc

@@ -25,12 +25,23 @@ namespace neuropod
 {
 namespace jni
 {
+// java_util_ArrayList is used to preset TesnorSpec data.
 jclass    java_util_ArrayList;
 jmethodID java_util_ArrayList_;
 jmethodID java_util_ArrayList_add;
 jmethodID java_util_ArrayList_get;
 jmethodID java_util_ArrayList_size;
 
+// The java_util_HashMap is needed by the infer method.
+// The infer method creates a Java hash map, which is converted from the output NeuropodValueMap.
+jclass    java_util_HashMap;
+jmethodID java_util_HashMap_;
+jmethodID java_util_HashMap_put;
+
+jclass    java_util_Map_Entry;
+jmethodID java_util_Map_Entry_getKey;
+jmethodID java_util_Map_Entry_getValue;
+
 jclass    com_uber_neuropod_TensorSpec;
 jmethodID com_uber_neuropod_TensorSpec_;
 
@@ -40,6 +51,10 @@ jmethodID com_uber_neuropod_Dimension_symbol_;
 
 jclass com_uber_neuropod_TensorType;
 
+jclass    com_uber_neuropod_NeuropodTensor;
+jmethodID com_uber_neuropod_NeuropodTensor_;
+jmethodID com_uber_neuropod_NeuropodTensor_getHandle;
+
 jclass com_uber_neuropod_NeuropodJNIException;
 
 jint JNI_VERSION = JNI_VERSION_1_8;
@@ -72,6 +87,15 @@ jint JNI_OnLoad(JavaVM *vm, void *reserved)
         java_util_ArrayList_get  = getMethodID(env, java_util_ArrayList, "get", "(I)Ljava/lang/Object;");
         java_util_ArrayList_size = getMethodID(env, java_util_ArrayList, "size", "()I");
 
+        java_util_HashMap  = static_cast<jclass>(env->NewGlobalRef(findClass(env, "java/util/HashMap")));
+        java_util_HashMap_ = getMethodID(env, java_util_HashMap, "<init>", "()V");
+        java_util_HashMap_put =
+            getMethodID(env, java_util_HashMap, "put", "(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;");
+
+        java_util_Map_Entry          = static_cast<jclass>(env->NewGlobalRef(findClass(env, "java/util/Map$Entry")));
+        java_util_Map_Entry_getKey   = getMethodID(env, java_util_Map_Entry, "getKey", "()Ljava/lang/Object;");
+        java_util_Map_Entry_getValue = getMethodID(env, java_util_Map_Entry, "getValue", "()Ljava/lang/Object;");
+
         com_uber_neuropod_TensorSpec =
             static_cast<jclass>(env->NewGlobalRef(findClass(env, "com/uber/neuropod/TensorSpec")));
         com_uber_neuropod_TensorSpec_ =
@@ -88,6 +112,12 @@ jint JNI_OnLoad(JavaVM *vm, void *reserved)
 
         com_uber_neuropod_TensorType =
             static_cast<jclass>(env->NewGlobalRef(findClass(env, "com/uber/neuropod/TensorType")));
+
+        com_uber_neuropod_NeuropodTensor =
+            static_cast<jclass>(env->NewGlobalRef(findClass(env, "com/uber/neuropod/NeuropodTensor")));
+        com_uber_neuropod_NeuropodTensor_ = getMethodID(env, com_uber_neuropod_NeuropodTensor, "<init>", "(J)V");
+        com_uber_neuropod_NeuropodTensor_getHandle =
+            getMethodID(env, com_uber_neuropod_NeuropodTensor, "getHandle", "()J");
     }
     catch (const std::exception &e)
     {
@@ -105,9 +135,12 @@ void JNI_OnUnload(JavaVM *vm, void *reserved)
     vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION);
     // Destroy the global references
     env->DeleteGlobalRef(java_util_ArrayList);
+    env->DeleteGlobalRef(java_util_HashMap);
+    env->DeleteGlobalRef(java_util_Map_Entry);
 
     env->DeleteGlobalRef(com_uber_neuropod_Dimension);
     env->DeleteGlobalRef(com_uber_neuropod_TensorSpec);
     env->DeleteGlobalRef(com_uber_neuropod_TensorType);
+    env->DeleteGlobalRef(com_uber_neuropod_NeuropodTensor);
     env->DeleteGlobalRef(com_uber_neuropod_NeuropodJNIException);
 }

source/neuropod/bindings/java/src/main/native/jclass_register.h

@@ -27,14 +27,26 @@ extern jmethodID java_util_ArrayList_add;
 extern jmethodID java_util_ArrayList_get;
 extern jmethodID java_util_ArrayList_size;
 
+extern jclass    java_util_HashMap;
+extern jmethodID java_util_HashMap_;
+extern jmethodID java_util_HashMap_put;
+
+extern jclass    java_util_Map_Entry;
+extern jmethodID java_util_Map_Entry_getKey;
+extern jmethodID java_util_Map_Entry_getValue;
+
+extern jclass com_uber_neuropod_TensorType;
+
 extern jclass    com_uber_neuropod_TensorSpec;
 extern jmethodID com_uber_neuropod_TensorSpec_;
 
 extern jclass    com_uber_neuropod_Dimension;
 extern jmethodID com_uber_neuropod_Dimension_value_;
 extern jmethodID com_uber_neuropod_Dimension_symbol_;
 
-extern jclass com_uber_neuropod_TensorType;
+extern jclass    com_uber_neuropod_NeuropodTensor;
+extern jmethodID com_uber_neuropod_NeuropodTensor_;
+extern jmethodID com_uber_neuropod_NeuropodTensor_getHandle;
 
 extern jclass com_uber_neuropod_NeuropodJNIException;