How to load custom op library when using c++?

Question

I have build a very simple custom op zero_out.dll with bazel, it works when using python.

import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.dll')
with tf.Session(''):
  zero_out_module.zero_out([[1, 2], [3, 4]]).eval()

However I have to run the inference using C++, is there any c++ api which has the similar function as tf.load_op_library, as it seems that a lot of registeration work has been done in tf.load_op_library, TF has not counterpart c++ API?

Answer 1

While there does not seem to be a public API for that in C++, the library loading functions are exposed in the TensorFlow API for C (which is the API that tf.load_library uses). There is no "nice" documentation for it, but you can find them in c/c_api.h:

// --------------------------------------------------------------------------
// Load plugins containing custom ops and kernels

// TF_Library holds information about dynamically loaded TensorFlow plugins.
typedef struct TF_Library TF_Library;

// Load the library specified by library_filename and register the ops and
// kernels present in that library.
//
// Pass "library_filename" to a platform-specific mechanism for dynamically
// loading a library. The rules for determining the exact location of the
// library are platform-specific and are not documented here.
//
// On success, place OK in status and return the newly created library handle.
// The caller owns the library handle.
//
// On failure, place an error status in status and return NULL.
TF_CAPI_EXPORT extern TF_Library* TF_LoadLibrary(const char* library_filename,
                                                 TF_Status* status);

// Get the OpList of OpDefs defined in the library pointed by lib_handle.
//
// Returns a TF_Buffer. The memory pointed to by the result is owned by
// lib_handle. The data in the buffer will be the serialized OpList proto for
// ops defined in the library.
TF_CAPI_EXPORT extern TF_Buffer TF_GetOpList(TF_Library* lib_handle);

// Frees the memory associated with the library handle.
// Does NOT unload the library.
TF_CAPI_EXPORT extern void TF_DeleteLibraryHandle(TF_Library* lib_handle);

These functions do actually call C++ code (see source in c/c_api.cc). However, the called functions, defined in core/framework/load_library.cc does not have a header to include. The workaround to use it in C++ code, which they use in c/c_api.cc, is to declare the function yourself, and link the TensorFlow library.

namespace tensorflow {
// Helpers for loading a TensorFlow plugin (a .so file).
Status LoadLibrary(const char* library_filename, void** result,
                   const void** buf, size_t* len);
}

As far as I can tell there is no API to unload the library. The C API allows you only to delete the library handle object. This done just by freeing the pointer, but if you want to avoid trouble you should probably use the freeing function given by TensorFlow, tensorflow::port:free, declared in core/platform/mem.h. Again, if you cannot not or don't want to include that, you can declare the function yourself and it should work as well.

namespace tensorflow {
namespace port {
void Free(void* ptr);
}
}