Refactor JNI code in C++ into Java code with JavaCPP #18

saudet · 2020-01-20T14:31:03Z

This is a WIP to remove the manually written JNI code from tensorflow-core-api. In the first commit for SavedModelBundle.java, I replaced about 100 lines of JNI code in C++ with about 40 lines of code in Java, and the unit tests still pass. We need to do the same with the rest of the JNI code, which shouldn't take me that long to do by myself, but if anyone is interested in helping, please let me know and I will give you write access to my fork!

There is also a lot of code in the Java section of the wrappers that is redundant with JavaCPP and that we will be able to remove in a second phase.

/cc @tzolov

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/SavedModelBundle.java

karllessard · 2020-01-20T22:02:05Z

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/SavedModelBundle.java

+        throw new IndexOutOfBoundsException("MetaGraphDef is too large to serialize into a byte[] array");
+      } else {
+        byte[] jmetagraph_def = new byte[(int)metagraph_def.length()];
+        new BytePointer(metagraph_def.data()).get(jmetagraph_def);


I suggest we move this array copy stunt in AbstractTF_Buffer as a utility method like toBytes(), since TF_Buffer is meant to be generic.

Alternatively, we could retain a reference to metaGraphDef and return its data as a ByteBuffer to the user instead of an array, as it is only used for deserializing the proto message... But I'm wondering why we just don't import those proto classes in the core API and return typed objects to the user like MetaGraphDef instead of leaving him the burden to serialize/deserialize the proto messages.

What do you think @sjamesr ? Is the original reason of leaving the protos outside the Java client was only to avoid an extra dependency to grpc?

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/SavedModelBundle.java

...core/tensorflow-core-api/src/main/java/org/tensorflow/internal/c_api/AbstractTF_Session.java

karllessard · 2020-01-22T13:31:49Z

@saudet , this PR is still marked as WIP, looks like it's ready to be merged now?

saudet · 2020-01-22T20:47:14Z

I was thinking we could do all of it together before merging? It shouldn't take me more than a few days to do...

karllessard · 2020-01-23T02:17:50Z

As you wish, we can do it iteratively with smaller PRs as well

saudet · 2020-01-23T03:27:36Z

Let's talk about it at the meeting tomorrow :)

karllessard

Looks good @saudet , I just left a few questions and remarks here and there.

I didn't compared yet the old JNI code with the new Java implementations but I think we can also rely on the Junit coverage for this.

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/EagerOperation.java

karllessard · 2020-01-29T02:00:49Z

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/EagerOperation.java

+      int length = TFE_OpGetOutputLength(handle, name, status);
+      status.throwExceptionIfNotOK();
+      return length;
+    }


Instead of starting new scopes in all of these methods, could it be simpler and more efficient to just create the status in a try-with-resource block when there are not other resource allocated?

try (TF_Status status = TF_Status.newStatus()) { ... }

Yes, it would be more efficient, but it would also make it more error-prone when we start creating other objects in there that may start doing temporary allocations.

Yeah... I would still prefer we go with the most efficient approach but it's up to you if you want to make the changes or not, we can merge it like this too.

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java

karllessard · 2020-01-29T02:43:01Z

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java

+        outputTensorHandles[i] = outputValues.get(TF_Tensor.class, i);
+      }
+
+      return runMetadata != null ? runMetadata.get() : null;


Question: did you ever ran a benchmark showing that all these context switches between the JVM and the native code (i.e. at each JavaCPP generated method) ends up to be as performant as when only one call was made (run in this case)?

Calls from the JVM to native code are very efficient, in the order of 30 ns. On the other hand, calls from native code to the JVM are typically very expensive, in the order of 300 ns, so it's almost certain that the new code here is going to be faster. I'll run a simple benchmark and post the results here just to confirm, but if you're worried about performance, we should think about writing a whole set of benchmarks to make sure there is never any regression in performance anywhere.

I've started to write some of them as well (like this one), I think with high-performance projects like TF we should have a set of benchmarks, not only it is useful to detect regression but also is a very good indicator of where we should spend time for optimization.

If you can add some, that would be great, but we can merge without if you want to create these later.

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Tensor.java

tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/TensorTest.java

saudet · 2020-01-29T05:29:40Z

I ran the initTensorByArrays benchmark from pull #20 (comment) with JMH's default settings and here is the kind of result I am getting on my machine:

Before (manually written JNI code):

Benchmark                        Mode  Cnt   Score   Error  Units
MyBenchmark.initTensorByArrays  thrpt   25  93.895 ± 1.580  ops/s

After (JNI refactored into Java with JavaCPP):

Benchmark                        Mode  Cnt    Score   Error  Units
MyBenchmark.initTensorByArrays  thrpt   25  110.598 ± 1.012  ops/s

So it does look like this refactoring also increases performance!

karllessard

So it does look like this refactoring also increases performance!

Great!

karllessard

All right, I'm merging this. There is a few unresolved conversation left that are not mandatory to be closed right now and could be part of another PR if needs be.

saudet requested a review from karllessard January 20, 2020 14:31

saudet changed the title ~~Refactor saved_model_bundle_jni into SavedModelBundle with JavaCPP~~ Refactor JNI code in C++ into Java code with JavaCPP Jan 20, 2020

karllessard reviewed Jan 21, 2020

View reviewed changes

saudet marked this pull request as ready for review January 27, 2020 22:27

saudet added 4 commits January 28, 2020 09:26

Refactor saved_model_bundle_jni into SavedModelBundle with JavaCPP

534e9d4

Fix code in response to review comments

e69874c

Refactor the rest of the JNI code into Java with JavaCPP

b8ff850

Fix formatting of switch statements and update URL in error message

55b7b5b

saudet force-pushed the remove-jni-code branch from dd6373b to 55b7b5b Compare January 28, 2020 03:35

karllessard requested changes Jan 29, 2020

View reviewed changes

saudet added 2 commits January 29, 2020 12:41

Remove calls to inefficient generic Pointer.put()

47125ef

Fix nits

1e22569

saudet requested a review from karllessard January 29, 2020 07:16

karllessard reviewed Jan 29, 2020

View reviewed changes

karllessard approved these changes Jan 29, 2020

View reviewed changes

karllessard merged commit b966d37 into tensorflow:master Jan 29, 2020

saudet mentioned this pull request Feb 5, 2020

Java maven snapshots and releases microsoft/onnxruntime#2675

Closed

saudet mentioned this pull request Jul 23, 2020

[RFC] MXNet 2.0 JVM Language development apache/mxnet#17783

Open

saudet mentioned this pull request Jan 5, 2021

Merging with JCuda and JOpenCL projects for better quality cuda interfaces bytedeco/javacpp-presets#475

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor JNI code in C++ into Java code with JavaCPP #18

Refactor JNI code in C++ into Java code with JavaCPP #18

saudet commented Jan 20, 2020

karllessard Jan 20, 2020

karllessard commented Jan 22, 2020

saudet commented Jan 22, 2020

karllessard commented Jan 23, 2020

saudet commented Jan 23, 2020

karllessard left a comment

karllessard Jan 29, 2020

saudet Jan 29, 2020 •

edited

Loading

karllessard Jan 29, 2020

karllessard Jan 29, 2020

saudet Jan 29, 2020 •

edited

Loading

karllessard Jan 29, 2020

saudet commented Jan 29, 2020

karllessard left a comment

karllessard left a comment

Refactor JNI code in C++ into Java code with JavaCPP #18

Refactor JNI code in C++ into Java code with JavaCPP #18

Conversation

saudet commented Jan 20, 2020

karllessard Jan 20, 2020

Choose a reason for hiding this comment

karllessard commented Jan 22, 2020

saudet commented Jan 22, 2020

karllessard commented Jan 23, 2020

saudet commented Jan 23, 2020

karllessard left a comment

Choose a reason for hiding this comment

karllessard Jan 29, 2020

Choose a reason for hiding this comment

saudet Jan 29, 2020 • edited Loading

Choose a reason for hiding this comment

karllessard Jan 29, 2020

Choose a reason for hiding this comment

karllessard Jan 29, 2020

Choose a reason for hiding this comment

saudet Jan 29, 2020 • edited Loading

Choose a reason for hiding this comment

karllessard Jan 29, 2020

Choose a reason for hiding this comment

saudet commented Jan 29, 2020

karllessard left a comment

Choose a reason for hiding this comment

karllessard left a comment

Choose a reason for hiding this comment

saudet Jan 29, 2020 •

edited

Loading

saudet Jan 29, 2020 •

edited

Loading