Java: Add a Graph class and the ability to import/export GraphDefs.

Also link in op definitions into libtensorflow-jni.so

Another step in the journey that is zplizzi#5.
Change: 141113176

tensorflow/java/BUILD

@@ -12,6 +12,35 @@ java_library(
     visibility = ["//visibility:public"],
 )
 
+java_test(
+    name = "GraphTest",
+    size = "small",
+    srcs = ["src/test/java/org/tensorflow/GraphTest.java"],
+    data = [":test_graph_def"],
+    test_class = "org.tensorflow.GraphTest",
+    deps = [
+        ":tensorflow",
+        "//external:junit",
+    ],
+)
+
+# Temporary targets to assist with unittests till the Java API becomes
+# expressive enough. Delete this target and corresponding python source code
+# when the Java API can construct graphs.
+genrule(
+    name = "test_graph_def",
+    outs = [":test_graph_def.data"],
+    cmd = "$(location :graphdef) $@",
+    tools = [":graphdef"],
+)
+
+py_binary(
+    name = "graphdef",
+    srcs = ["src/test/python/graphdef.py"],
+    srcs_version = "PY2AND3",
+    deps = ["//tensorflow:tensorflow_py"],
+)
+
 java_test(
     name = "TensorFlowTest",
     size = "small",

tensorflow/java/src/main/java/org/tensorflow/Graph.java

@@ -0,0 +1,105 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+package org.tensorflow;
+
+/**
+ * A data flow graph representing a TensorFlow computation.
+ *
+ * <p>Instances of a Graph are thread-safe.
+ *
+ * <p><b>WARNING:</b> Resources consumed by the Graph object msut be explicitly freed by invoking
+ * the {@link #close()} method then the Graph object is no longer needed.
+ */
+public final class Graph implements AutoCloseable {
+
+  /** Create an empty Graph. */
+  public Graph() {
+    nativeHandle = allocate();
+  }
+
+  /**
+   * Release resources associated with the Graph.
+   *
+   * <p>A Graph is no longer usable after close returns.
+   */
+  @Override
+  public void close() {
+    synchronized (nativeHandleLock) {
+      if (nativeHandle != 0) {
+        delete(nativeHandle);
+        nativeHandle = 0;
+      }
+    }
+  }
+
+  /**
+   * Import a serialized representation of a TensorFlow graph.
+   *
+   * <p>The serialized representation of the graph, often referred to as a <i>GraphDef</i>, can be
+   * generated by {@link #toGraphDef()} and equivalents in other language APIs.
+   *
+   * @throws IllegalArgumentException if graphDef is not a recognized serialization of a graph.
+   * @see #importGraphDef(byte[], String)
+   */
+  public void importGraphDef(byte[] graphDef) throws IllegalArgumentException {
+    importGraphDef(graphDef, "");
+  }
+
+  /**
+   * Import a serialized representation of a TensorFlow graph.
+   *
+   * @param graphDef the serialized representation of a TensorFlow graph.
+   * @param prefix a prefix that will be prepended to names in graphDef
+   * @throws IllegalArgumentException if graphDef is not a recognized serialization of a graph.
+   * @see #importGraphDef(byte[])
+   */
+  public void importGraphDef(byte[] graphDef, String prefix) throws IllegalArgumentException {
+    if (graphDef == null || prefix == null) {
+      throw new IllegalArgumentException("graphDef and prefix cannot be null");
+    }
+    synchronized (nativeHandleLock) {
+      importGraphDef(nativeHandle, graphDef, prefix);
+    }
+  }
+
+  /**
+   * Generate a serialized representation of the Graph.
+   *
+   * @see #importGraphDef(byte[])
+   * @see #importGraphDef(byte[], String)
+   */
+  public byte[] toGraphDef() {
+    synchronized (nativeHandleLock) {
+      return toGraphDef(nativeHandle);
+    }
+  }
+
+  private final Object nativeHandleLock = new Object();
+  private long nativeHandle;
+
+  private static native long allocate();
+
+  private static native void delete(long handle);
+
+  private static native void importGraphDef(long handle, byte[] graphDef, String prefix)
+      throws IllegalArgumentException;
+
+  private static native byte[] toGraphDef(long handle);
+
+  static {
+    TensorFlow.init();
+  }
+}

tensorflow/java/src/main/native/BUILD

@@ -18,6 +18,7 @@ tf_cuda_library(
     includes = ["."],
     deps = [
         "//tensorflow/c:c_api",
+        "//tensorflow/core:ops",
     ],
     alwayslink = 1,
 )

tensorflow/java/src/main/native/exception_jni.cc

@@ -21,6 +21,7 @@ limitations under the License.
 const char kIllegalArgumentException[] = "java/lang/IllegalArgumentException";
 const char kIllegalStateException[] = "java/lang/IllegalStateException";
 const char kNullPointerException[] = "java/lang/NullPointerException";
+const char kIndexOutOfBoundsException[] = "java/lang/IndexOutOfBoundsException";
 
 void throwException(JNIEnv* env, const char* clazz, const char* fmt, ...) {
   va_list args;

tensorflow/java/src/main/native/exception_jni.h

@@ -25,6 +25,7 @@ extern "C" {
 extern const char kIllegalArgumentException[];
 extern const char kIllegalStateException[];
 extern const char kNullPointerException[];
+extern const char kIndexOutOfBoundsException[];
 
 void throwException(JNIEnv* env, const char* clazz, const char* fmt, ...);
 

tensorflow/java/src/main/native/graph_jni.cc

@@ -0,0 +1,106 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/java/src/main/native/graph_jni.h"
+
+#include <limits>
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/java/src/main/native/exception_jni.h"
+
+namespace {
+TF_Graph* requireHandle(JNIEnv* env, jlong handle) {
+  if (handle == 0) {
+    throwException(env, kNullPointerException,
+                   "close() has been called on the Graph");
+    return nullptr;
+  }
+  return reinterpret_cast<TF_Graph*>(handle);
+}
+}  // namespace
+
+JNIEXPORT jlong JNICALL Java_org_tensorflow_Graph_allocate(JNIEnv*, jclass) {
+  return reinterpret_cast<jlong>(TF_NewGraph());
+}
+
+JNIEXPORT void JNICALL Java_org_tensorflow_Graph_delete(JNIEnv*, jclass,
+                                                        jlong handle) {
+  if (handle == 0) return;
+  TF_DeleteGraph(reinterpret_cast<TF_Graph*>(handle));
+}
+
+JNIEXPORT void JNICALL Java_org_tensorflow_Graph_importGraphDef(
+    JNIEnv* env, jclass clazz, jlong handle, jbyteArray graph_def,
+    jstring prefix) {
+  TF_Graph* g = requireHandle(env, handle);
+  if (g == nullptr) return;
+
+  TF_ImportGraphDefOptions* opts = TF_NewImportGraphDefOptions();
+
+  jboolean is_copy;
+  const char* cprefix = env->GetStringUTFChars(prefix, &is_copy);
+  TF_ImportGraphDefOptionsSetPrefix(opts, cprefix);
+  env->ReleaseStringUTFChars(prefix, cprefix);
+
+  static_assert(sizeof(jbyte) == 1, "unexpected size of the jbyte type");
+  jbyte* bytes = env->GetByteArrayElements(graph_def, &is_copy);
+  TF_Buffer* buf =
+      TF_NewBufferFromString(bytes, env->GetArrayLength(graph_def));
+  TF_Status* status = TF_NewStatus();
+
+  TF_GraphImportGraphDef(g, buf, opts, status);
+  if (TF_GetCode(status) != TF_OK) {
+    // TODO(ashankar): Throw exceptions that will distinguish between different
+    // status codes.
+    throwException(env, kIllegalArgumentException, TF_Message(status));
+    // Do not return here, many resources need to be cleaned up.
+  }
+
+  TF_DeleteStatus(status);
+  TF_DeleteBuffer(buf);
+  env->ReleaseByteArrayElements(graph_def, bytes, JNI_ABORT);
+
+  TF_DeleteImportGraphDefOptions(opts);
+}
+
+JNIEXPORT jbyteArray JNICALL
+Java_org_tensorflow_Graph_toGraphDef(JNIEnv* env, jclass clazz, jlong handle) {
+  jbyteArray ret = nullptr;
+  TF_Graph* g = requireHandle(env, handle);
+  if (g == nullptr) return ret;
+
+  TF_Buffer* buf = TF_NewBuffer();
+  TF_Status* status = TF_NewStatus();
+  TF_GraphToGraphDef(g, buf, status);
+  if (TF_GetCode(status) != TF_OK) {
+    // TODO(ashankar): Throw exceptions that will distinguish between different
+    // status codes.
+    throwException(env, kIllegalStateException, TF_Message(status));
+  } else {
+    // sizeof(jsize) is less than sizeof(size_t) on some platforms.
+    if (buf->length > std::numeric_limits<jint>::max()) {
+      throwException(env, kIndexOutOfBoundsException,
+                     "GraphDef is too large to serialize into a byte[] array");
+    } else {
+      static_assert(sizeof(jbyte) == 1, "unexpected size of the jbyte type");
+      jint ret_len = static_cast<jint>(buf->length);
+      ret = env->NewByteArray(ret_len);
+      env->SetByteArrayRegion(ret, 0, ret_len,
+                              static_cast<const jbyte*>(buf->data));
+    }
+  }
+  TF_DeleteStatus(status);
+  TF_DeleteBuffer(buf);
+  return ret;
+}

tensorflow/java/src/main/native/graph_jni.h

@@ -0,0 +1,62 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_JAVA_GRAPH_JNI_H_
+#define TENSORFLOW_JAVA_GRAPH_JNI_H_
+
+#include <jni.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/*
+ * Class:     org_tensorflow_Graph
+ * Method:    allocate
+ * Signature: ()J
+ */
+JNIEXPORT jlong JNICALL Java_org_tensorflow_Graph_allocate(JNIEnv *, jclass);
+
+/*
+ * Class:     org_tensorflow_Graph
+ * Method:    delete
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_org_tensorflow_Graph_delete(JNIEnv *, jclass,
+                                                        jlong);
+
+/*
+ * Class:     org_tensorflow_Graph
+ * Method:    importGraphDef
+ * Signature: (J[BLjava/lang/String;)V
+ */
+JNIEXPORT void JNICALL Java_org_tensorflow_Graph_importGraphDef(JNIEnv *,
+                                                                jclass, jlong,
+                                                                jbyteArray,
+                                                                jstring);
+
+/*
+ * Class:     org_tensorflow_Graph
+ * Method:    toGraphDef
+ * Signature: (J)[B
+ */
+JNIEXPORT jbyteArray JNICALL Java_org_tensorflow_Graph_toGraphDef(JNIEnv *,
+                                                                  jclass,
+                                                                  jlong);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif  // __cplusplus
+#endif  // TENSORFLOW_JAVA_GRAPH_JNI_H_

tensorflow/java/src/test/java/org/tensorflow/GraphTest.java

@@ -0,0 +1,63 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+package org.tensorflow;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link org.tensorflow.Graph}. */
+@RunWith(JUnit4.class)
+public class GraphTest {
+
+  @Test
+  public void graphDefImportAndExport() {
+    try (Graph g = new Graph()) {
+      final byte[] inGraphDef = Files.readAllBytes(Paths.get("tensorflow/java/test_graph_def.data"));
+      g.importGraphDef(inGraphDef);
+      final byte[] outGraphDef = g.toGraphDef();
+      // The graphs may not be identical as the proto format allows the same message
+      // to be encoded in multiple ways. Once the Graph API is expressive enough
+      // to construct graphs and query for nodes/operations, use that.
+      // Till then a very crude test:
+      assertEquals(inGraphDef.length, outGraphDef.length);
+    } catch (Exception e) {
+      fail("Unexpected exception: " + e);
+    }
+  }
+
+  @Test
+  public void failImportOnInvalidGraphDefs() {
+    try (Graph g = new Graph()) {
+      try {
+        g.importGraphDef(null);
+      } catch (IllegalArgumentException e) {
+        // expected exception.
+      }
+
+      try {
+        g.importGraphDef(new byte[] {1});
+      } catch (IllegalArgumentException e) {
+        // expected exception.
+      }
+    }
+  }
+}