• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2016 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# ==============================================================================
15r"""Computes a header file to be used with SELECTIVE_REGISTRATION.
16
17See the executable wrapper, print_selective_registration_header.py, for more
18information.
19"""
20
21from __future__ import absolute_import
22from __future__ import division
23from __future__ import print_function
24
25import json
26import os
27import sys
28
29from google.protobuf import text_format
30from tensorflow.core.framework import graph_pb2
31from tensorflow.python.platform import gfile
32from tensorflow.python.platform import tf_logging
33from tensorflow.python.util import _pywrap_kernel_registry
34
35# Usually, we use each graph node to induce registration of an op and
36# corresponding kernel; nodes without a corresponding kernel (perhaps due to
37# attr types) generate a warning but are otherwise ignored. Ops in this set are
38# registered even if there's no corresponding kernel.
39OPS_WITHOUT_KERNEL_ALLOWLIST = frozenset([
40    # AccumulateNV2 is rewritten away by AccumulateNV2RemovePass; see
41    # core/common_runtime/accumulate_n_optimizer.cc.
42    'AccumulateNV2'
43])
44FLEX_PREFIX = b'Flex'
45FLEX_PREFIX_LENGTH = len(FLEX_PREFIX)
46
47
48def _get_ops_from_ops_list(input_file):
49  """Gets the ops and kernels needed from the ops list file."""
50  ops = set()
51  ops_list_str = gfile.GFile(input_file, 'r').read()
52  if not ops_list_str:
53    raise Exception('Input file should not be empty')
54  ops_list = json.loads(ops_list_str)
55  for op, kernel in ops_list:
56    op_and_kernel = (op, kernel if kernel else None)
57    ops.add(op_and_kernel)
58  return ops
59
60
61def _get_ops_from_graphdef(graph_def):
62  """Gets the ops and kernels needed from the tensorflow model."""
63  ops = set()
64  for node_def in graph_def.node:
65    if not node_def.device:
66      node_def.device = '/cpu:0'
67    kernel_class = _pywrap_kernel_registry.TryFindKernelClass(
68        node_def.SerializeToString())
69    op = str(node_def.op)
70    if kernel_class or op in OPS_WITHOUT_KERNEL_ALLOWLIST:
71      op_and_kernel = (op, str(kernel_class.decode('utf-8'))
72                       if kernel_class else None)
73      ops.add(op_and_kernel)
74    else:
75      print('Warning: no kernel found for op %s' % node_def.op, file=sys.stderr)
76  return ops
77
78
79def get_ops_and_kernels(proto_fileformat, proto_files, default_ops_str):
80  """Gets the ops and kernels needed from the model files."""
81  ops = set()
82
83  for proto_file in proto_files:
84    tf_logging.info('Loading proto file %s', proto_file)
85    # Load ops list file.
86    if proto_fileformat == 'ops_list':
87      ops = ops.union(_get_ops_from_ops_list(proto_file))
88      continue
89
90    # Load GraphDef.
91    file_data = gfile.GFile(proto_file, 'rb').read()
92    if proto_fileformat == 'rawproto':
93      graph_def = graph_pb2.GraphDef.FromString(file_data)
94    else:
95      assert proto_fileformat == 'textproto'
96      graph_def = text_format.Parse(file_data, graph_pb2.GraphDef())
97    ops = ops.union(_get_ops_from_graphdef(graph_def))
98
99  # Add default ops.
100  if default_ops_str and default_ops_str != 'all':
101    for s in default_ops_str.split(','):
102      op, kernel = s.split(':')
103      op_and_kernel = (op, kernel)
104      if op_and_kernel not in ops:
105        ops.add(op_and_kernel)
106
107  return list(sorted(ops))
108
109
110def get_header_from_ops_and_kernels(ops_and_kernels,
111                                    include_all_ops_and_kernels):
112  """Returns a header for use with tensorflow SELECTIVE_REGISTRATION.
113
114  Args:
115    ops_and_kernels: a set of (op_name, kernel_class_name) pairs to include.
116    include_all_ops_and_kernels: if True, ops_and_kernels is ignored and all op
117      kernels are included.
118
119  Returns:
120    the string of the header that should be written as ops_to_register.h.
121  """
122  ops = set(op for op, _ in ops_and_kernels)
123  result_list = []
124
125  def append(s):
126    result_list.append(s)
127
128  _, script_name = os.path.split(sys.argv[0])
129  append('// This file was autogenerated by %s' % script_name)
130  append('#ifndef OPS_TO_REGISTER')
131  append('#define OPS_TO_REGISTER')
132
133  if include_all_ops_and_kernels:
134    append('#define SHOULD_REGISTER_OP(op) true')
135    append('#define SHOULD_REGISTER_OP_KERNEL(clz) true')
136    append('#define SHOULD_REGISTER_OP_GRADIENT true')
137  else:
138    line = """
139    namespace {
140      constexpr const char* skip(const char* x) {
141        return (*x) ? (*x == ' ' ? skip(x + 1) : x) : x;
142      }
143
144      constexpr bool isequal(const char* x, const char* y) {
145        return (*skip(x) && *skip(y))
146                   ? (*skip(x) == *skip(y) && isequal(skip(x) + 1, skip(y) + 1))
147                   : (!*skip(x) && !*skip(y));
148      }
149
150      template<int N>
151      struct find_in {
152        static constexpr bool f(const char* x, const char* const y[N]) {
153          return isequal(x, y[0]) || find_in<N - 1>::f(x, y + 1);
154        }
155      };
156
157      template<>
158      struct find_in<0> {
159        static constexpr bool f(const char* x, const char* const y[]) {
160          return false;
161        }
162      };
163    }  // end namespace
164    """
165    line += 'constexpr const char* kNecessaryOpKernelClasses[] = {\n'
166    for _, kernel_class in ops_and_kernels:
167      if kernel_class is None:
168        continue
169      line += '"%s",\n' % kernel_class
170    line += '};'
171    append(line)
172    append('#define SHOULD_REGISTER_OP_KERNEL(clz) '
173           '(find_in<sizeof(kNecessaryOpKernelClasses) '
174           '/ sizeof(*kNecessaryOpKernelClasses)>::f(clz, '
175           'kNecessaryOpKernelClasses))')
176    append('')
177
178    append('constexpr inline bool ShouldRegisterOp(const char op[]) {')
179    append('  return false')
180    for op in sorted(ops):
181      append('     || isequal(op, "%s")' % op)
182    append('  ;')
183    append('}')
184    append('#define SHOULD_REGISTER_OP(op) ShouldRegisterOp(op)')
185    append('')
186
187    append('#define SHOULD_REGISTER_OP_GRADIENT ' +
188           ('true' if 'SymbolicGradient' in ops else 'false'))
189
190  append('#endif')
191  return '\n'.join(result_list)
192
193
194def get_header(graphs,
195               proto_fileformat='rawproto',
196               default_ops='NoOp:NoOp,_Recv:RecvOp,_Send:SendOp'):
197  """Computes a header for use with tensorflow SELECTIVE_REGISTRATION.
198
199  Args:
200    graphs: a list of paths to GraphDef files to include.
201    proto_fileformat: optional format of proto file, either 'textproto',
202      'rawproto' (default) or ops_list. The ops_list is the file contain the
203      list of ops in JSON format, Ex: "[["Transpose", "TransposeCpuOp"]]".
204    default_ops: optional comma-separated string of operator:kernel pairs to
205      always include implementation for. Pass 'all' to have all operators and
206      kernels included. Default: 'NoOp:NoOp,_Recv:RecvOp,_Send:SendOp'.
207
208  Returns:
209    the string of the header that should be written as ops_to_register.h.
210  """
211  ops_and_kernels = get_ops_and_kernels(proto_fileformat, graphs, default_ops)
212  if not ops_and_kernels:
213    print('Error reading graph!')
214    return 1
215
216  return get_header_from_ops_and_kernels(ops_and_kernels, default_ops == 'all')
217