• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 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"""mlir is an experimental library that provides support APIs for MLIR."""
16
17from tensorflow.python import pywrap_mlir
18from tensorflow.python.util.tf_export import tf_export
19
20
21@tf_export('mlir.experimental.convert_graph_def')
22def convert_graph_def(graph_def,
23                      pass_pipeline='tf-standard-pipeline',
24                      show_debug_info=False):
25  """Import a GraphDef and convert it to a textual MLIR module.
26
27  This API is only intended for inspecting the internals of TensorFlow and the
28  string returned is at the moment intended for debugging purposes.
29
30  Args:
31    graph_def: An object of type graph_pb2.GraphDef or a textual proto
32      representation of a valid GraphDef.
33    pass_pipeline: A textual description of an MLIR Pass Pipeline to run on the
34      module, see MLIR documentation for the
35      [textual pass pipeline syntax](https://mlir.llvm.org/docs/PassManagement/#textual-pass-pipeline-specification).
36    show_debug_info: Whether to include locations in the emitted textual form.
37
38  Returns:
39    A textual representation of the MLIR module corresponding to the graphdef.
40
41  Raises:
42    InvalidArgumentError: if graph_def is invalid or cannot be converted to
43      MLIR.
44
45  """
46  return pywrap_mlir.import_graphdef(graph_def, pass_pipeline, show_debug_info)
47
48
49@tf_export('mlir.experimental.convert_function')
50def convert_function(concrete_function,
51                     pass_pipeline='tf-standard-pipeline',
52                     show_debug_info=False):
53  """Import a ConcreteFunction and convert it to a textual MLIR module.
54
55  This API is only intended for inspecting the internals of TensorFlow and the
56  string returned is at the moment intended for debugging purposes.
57
58  A [tf.function](https://www.tensorflow.org/api_docs/python/tf/function) can be
59  imported and converted from TensorFlow to TensorFlow MLIR with this API by
60  extracting its ConcreteFunction (eagerly-executing wrapper around a
61  [tf.Graph](https://www.tensorflow.org/api_docs/python/tf/Graph)).
62
63  For example:
64  >>> @tf.function
65  ... def add(a, b):
66  ...   return a + b
67
68  >>> concrete_function = add.get_concrete_function(
69  ...     tf.TensorSpec(None, tf.dtypes.float32),
70  ...     tf.TensorSpec(None, tf.dtypes.float32))
71  >>> tf.mlir.experimental.convert_function(concrete_function)
72  '...module attributes {...} {...}...'
73
74  Args:
75    concrete_function: An object of type ConcreteFunction.
76    pass_pipeline: A textual description of an MLIR Pass Pipeline to run on the
77      module, see MLIR documentation for the
78      [textual pass pipeline syntax](https://mlir.llvm.org/docs/PassManagement/#textual-pass-pipeline-specification).
79    show_debug_info: Whether to include locations in the emitted textual form.
80
81  Returns:
82    A textual representation of the MLIR module corresponding to the
83    ConcreteFunction.
84
85  Raises:
86    InvalidArgumentError: if concrete_function is invalid or cannot be converted
87      to MLIR.
88
89  """
90  return pywrap_mlir.import_function(concrete_function, pass_pipeline,
91                                     show_debug_info)
92