1# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15
16"""Function for loading TensorFlow plugins."""
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import hashlib
22import imp
23import sys
24import threading  # pylint: disable=unused-import
25
26from tensorflow.core.framework import op_def_pb2
27from tensorflow.core.lib.core import error_codes_pb2  # pylint: disable=unused-import
28from tensorflow.python import pywrap_tensorflow as py_tf
29from tensorflow.python.framework import errors_impl
30from tensorflow.python.util import compat
31from tensorflow.python.util.tf_export import tf_export
32
33
34@tf_export('load_op_library')
35def load_op_library(library_filename):
36  """Loads a TensorFlow plugin, containing custom ops and kernels.
37
38  Pass "library_filename" to a platform-specific mechanism for dynamically
39  loading a library. The rules for determining the exact location of the
40  library are platform-specific and are not documented here. When the
41  library is loaded, ops and kernels registered in the library via the
42  `REGISTER_*` macros are made available in the TensorFlow process. Note
43  that ops with the same name as an existing op are rejected and not
44  registered with the process.
45
46  Args:
47    library_filename: Path to the plugin.
48      Relative or absolute filesystem path to a dynamic library file.
49
50  Returns:
51    A python module containing the Python wrappers for Ops defined in
52    the plugin.
53
54  Raises:
55    RuntimeError: when unable to load the library or get the python wrappers.
56  """
57  with errors_impl.raise_exception_on_not_ok_status() as status:
58    lib_handle = py_tf.TF_LoadLibrary(library_filename, status)
59
60  op_list_str = py_tf.TF_GetOpList(lib_handle)
61  op_list = op_def_pb2.OpList()
62  op_list.ParseFromString(compat.as_bytes(op_list_str))
63  wrappers = py_tf.GetPythonWrappers(op_list_str)
64
65  # Delete the library handle to release any memory held in C
66  # that are no longer needed.
67  py_tf.TF_DeleteLibraryHandle(lib_handle)
68
69  # Get a unique name for the module.
70  module_name = hashlib.md5(wrappers).hexdigest()
71  if module_name in sys.modules:
72    return sys.modules[module_name]
73  module = imp.new_module(module_name)
74  # pylint: disable=exec-used
75  exec(wrappers, module.__dict__)
76  # Stash away the library handle for making calls into the dynamic library.
77  module.LIB_HANDLE = lib_handle
78  # OpDefs of the list of ops defined in the library.
79  module.OP_LIST = op_list
80  sys.modules[module_name] = module
81  return module
82
83
84@tf_export('load_file_system_library')
85def load_file_system_library(library_filename):
86  """Loads a TensorFlow plugin, containing file system implementation.
87
88  Pass `library_filename` to a platform-specific mechanism for dynamically
89  loading a library. The rules for determining the exact location of the
90  library are platform-specific and are not documented here.
91
92  Args:
93    library_filename: Path to the plugin.
94      Relative or absolute filesystem path to a dynamic library file.
95
96  Returns:
97    None.
98
99  Raises:
100    RuntimeError: when unable to load the library.
101  """
102  with errors_impl.raise_exception_on_not_ok_status() as status:
103    lib_handle = py_tf.TF_LoadLibrary(library_filename, status)
104