Understand kernel registration and binding in TensorFlow

Question

I am new to TensorFlow and am looking at custom op development right now. I already read the official tutorial, but I feel that a lot is going on behind the scenes, and I don't always want to put my user statements in the user_ops directory.

So I reviewed the word2vec example

which uses the custom "Skipgram" op, the registration of which is defined here:
/word2vec_ops.cc
and the kernel implementation of which is here:
/word2vec_kernels.cc

By looking at the build file, I tried to create separate targets

1) bazel build -c opt tensorflow/models/embedding:word2vec_ops
This generates a bunch of object files as expected.

2) bazel build -c opt tensorflow/models/embedding:word2vec_kernels
The same goes for that.

3) bazel build -c opt tensorflow/models/embedding:word2vec_kernels:gen_word2vec

This latest build uses a custom rule, namely tf_op_gen_wrapper_py https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorflow.bzl#L197-L231

It is interesting to note that this depends only on Op Registration, and not on the kernel itself.

In the end, if I build py_binary using

bazel build -c opt tensorflow/models/embedding:word2vec

It works fine, but I don’t see where and how the C ++ kernel code is related?

In addition, I would also like to understand the tf_op_gen_wrapper_py rule and the whole compilation / linking procedure that goes behind the scenes for registering ops.

Thanks.

Answer 1

When adding a new kind of operation to TensorFlow , there are two main steps:

• Register "op" , which includes an interface definition for the operation and

• Registration of one or more "cores" , which includes the definition of implementation (s) for the operation, possibly with specialized implementations for different data types or device types (for example, CPU or GPU).

Both steps include writing C ++ code. Registration op uses the REGISTER_OP() macro, and the kernel registration uses the REGISTER_KERNEL_BUILDER() macro . These macros create static initializers that run when the module containing them loads. There are two main mechanisms for registering operations and the kernel:

• Static binding to the TensorFlow core library and static initialization.

• Dynamic binding at runtime using the tf.load_op_library() function.

In the case of "Skipgram" we use parameter 1 (static binding). The operations are tied to the TensorFlow core library here , and the cores are linked in. (Note that this is not ideal: the word2vec statements were created before we had tf.load_op_library() , and therefore there was no mechanism for dynamically linking them). Therefore, operators and kernels are logged the first time TensorFlow is loaded (in import tensorflow as tf ). If they were created today, they would be dynamically loaded, so that they would only be registered if necessary. (SyntaxNet code has an example of dynamic loading.)

tf_op_gen_wrapper_py() rule in Bazel takes a list of dependencies between operating libraries and generates Python wrappers for these ops. The reason this rule only depends on op registration is because Python wrappers are completely defined by the op interface, which is defined in the op registry. Notably, the Python interface does not know if specialized kernels exist for a particular type or device. A cover generator associates op registration with a simple C ++ binary that generates Python code for each op registered. Please note: if you use tf.load_op_library() , you do not need to call the shell generator yourself, because tf.load_op_library() generates the necessary code at runtime.