chore: import upstream snapshot with attribution

tensorflow/python/compiler/mlir/BUILD

@@ -0,0 +1,45 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "mlir",
+    srcs = ["mlir.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python:pywrap_mlir",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_test(
+    name = "mlir_test",
+    srcs = ["mlir_test.py"],
+    data = [
+        "multi_add.tflite",
+    ],
+    strict_deps = True,
+    tags = [
+        "no_mac",  # TODO(b/291995580): Reenable.
+        "no_pip",
+    ],
+    deps = [
+        ":mlir",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow/python:pywrap_mlir",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:logging_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:resource_loader",
+    ],
+)

tensorflow/python/compiler/mlir/mlir.py

@@ -0,0 +1,206 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""mlir is an experimental library that provides support APIs for MLIR."""
+
+from tensorflow.python import pywrap_mlir
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export('mlir.experimental.convert_graph_def')
+def convert_graph_def(
+    graph_def, pass_pipeline='tf-standard-pipeline', show_debug_info=False
+):
+  """Import a GraphDef and convert it to a textual MLIR module.
+
+  This API is only intended for inspecting the internals of TensorFlow and the
+  string returned is at the moment intended for debugging purposes.
+
+  Args:
+    graph_def: An object of type graph_pb2.GraphDef or a textual proto
+      representation of a valid GraphDef.
+    pass_pipeline: A textual description of an MLIR Pass Pipeline to run on the
+      module, see MLIR documentation for the [textual pass pipeline
+      syntax](https://mlir.llvm.org/docs/PassManagement/#textual-pass-pipeline-specification).
+    show_debug_info: Whether to include locations in the emitted textual form.
+
+  Returns:
+    A textual representation of the MLIR module corresponding to the graphdef.
+
+  Raises:
+    InvalidArgumentError: if graph_def is invalid or cannot be converted to
+      MLIR.
+  """
+  return pywrap_mlir.import_graphdef(graph_def, pass_pipeline, show_debug_info)
+
+
+@tf_export('mlir.experimental.convert_function')
+def convert_function(
+    concrete_function,
+    pass_pipeline='tf-standard-pipeline',
+    show_debug_info=False,
+):
+  """Import a ConcreteFunction and convert it to a textual MLIR module.
+
+  This API is only intended for inspecting the internals of TensorFlow and the
+  string returned is at the moment intended for debugging purposes.
+
+  A [tf.function](https://www.tensorflow.org/api_docs/python/tf/function) can be
+  imported and converted from TensorFlow to TensorFlow MLIR with this API by
+  extracting its ConcreteFunction (eagerly-executing wrapper around a
+  [tf.Graph](https://www.tensorflow.org/api_docs/python/tf/Graph)).
+
+  For example:
+  >>> @tf.function
+  ... def add(a, b):
+  ...   return a + b
+
+  >>> concrete_function = add.get_concrete_function(
+  ...     tf.TensorSpec(None, tf.dtypes.float32),
+  ...     tf.TensorSpec(None, tf.dtypes.float32))
+  >>> tf.mlir.experimental.convert_function(concrete_function)
+  '...module attributes {...} {...}...'
+
+  Args:
+    concrete_function: An object of type ConcreteFunction.
+    pass_pipeline: A textual description of an MLIR Pass Pipeline to run on the
+      module, see MLIR documentation for the [textual pass pipeline
+      syntax](https://mlir.llvm.org/docs/PassManagement/#textual-pass-pipeline-specification).
+    show_debug_info: Whether to include locations in the emitted textual form.
+
+  Returns:
+    A textual representation of the MLIR module corresponding to the
+    ConcreteFunction.
+
+  Raises:
+    InvalidArgumentError: if concrete_function is invalid or cannot be converted
+      to MLIR.
+  """
+  return pywrap_mlir.import_function(
+      concrete_function, pass_pipeline, show_debug_info
+  )
+
+
+@tf_export('mlir.experimental.convert_saved_model')
+def convert_saved_model(
+    saved_model_path, exported_names, show_debug_info=False
+):
+  """Converts a SavedModel to MLIR module.
+
+  Args:
+    saved_model_path: Path to SavedModel.
+    exported_names: Names to export.
+    show_debug_info: Whether to include locations in the emitted textual form.
+
+  Returns:
+    A textual representation of the MLIR module corresponding to the
+    SavedModel.
+  """
+  return pywrap_mlir.experimental_convert_saved_model_to_mlir(
+      saved_model_path, exported_names, show_debug_info
+  )
+
+
+@tf_export('mlir.experimental.convert_saved_model_v1')
+def convert_saved_model_v1(
+    saved_model_path,
+    exported_names,
+    tags,
+    lift_variables,
+    include_variables_in_initializers,
+    upgrade_legacy=True,
+    show_debug_info=False,
+):
+  """Converts a v1 SavedModel to MLIR module.
+
+  Args:
+    saved_model_path: Path to SavedModel.
+    exported_names: Names to export.
+    tags: MetaGraphDef to be loaded is identified by the supplied tags.
+    lift_variables: Whether to promote tf.VarHandleOp to resource arguments.
+    include_variables_in_initializers: Keeps the variables in initializers
+      before lifting variables.
+    upgrade_legacy: Functionalize the input graph before importing.
+    show_debug_info: Whether to include locations in the emitted textual form.
+
+  Returns:
+    A textual representation of the MLIR module corresponding to the
+    SavedModule.
+  """
+  return pywrap_mlir.experimental_convert_saved_model_v1_to_mlir(
+      saved_model_path,
+      exported_names,
+      tags,
+      lift_variables,
+      include_variables_in_initializers,
+      upgrade_legacy,
+      show_debug_info,
+  )
+
+
+@tf_export('mlir.experimental.run_pass_pipeline')
+def run_pass_pipeline(mlir_txt, pass_pipeline, show_debug_info=False):
+  """Runs a pipeline over input module.
+
+  Args:
+    mlir_txt: Textual representation of the MLIR module.
+    pass_pipeline: Pass pipeline to run on module.
+    show_debug_info: Whether to include locations in the emitted textual form.
+
+  Returns:
+    A textual representation of the MLIR module corresponding to the
+    transformed module.
+  """
+  return pywrap_mlir.experimental_run_pass_pipeline(
+      mlir_txt, pass_pipeline, show_debug_info
+  )
+
+
+@tf_export('mlir.experimental.write_bytecode')
+def experimental_write_bytecode(filename, mlir_txt):
+  """Writes an MLIR module out as bytecode.
+
+  Args:
+    filename: The filename to write to.
+    mlir_txt: The MLIR module in textual format.
+  """
+  pywrap_mlir.experimental_write_bytecode(filename, mlir_txt)
+
+
+@tf_export('mlir.experimental.tflite_to_tosa_bytecode')
+def tflite_to_tosa_bytecode(
+    flatbuffer,
+    bytecode,
+    use_external_constant=False,
+    ordered_input_arrays=None,
+    ordered_output_arrays=None,
+):
+  """Converts TFLite flatbuffer to TOSA dialect in MLIR bytecode.
+
+  Args:
+    flatbuffer: Path to flatbuffer.
+    bytecode: Path to output bytecode.
+    use_external_constant: Whether to create `tfl.external_const` instead of
+      `tfl.const`.
+    ordered_input_arrays:
+    ordered_output_arrays: If ordered_output_arrays is not empty, then the
+      function will only return nodes in ordered_output_arrays in the same order
+  """
+  pywrap_mlir.experimental_tflite_to_tosa_bytecode(
+      flatbuffer,
+      bytecode,
+      use_external_constant,
+      ordered_input_arrays,
+      ordered_output_arrays,
+  )

tensorflow/python/compiler/mlir/mlir_test.py

@@ -0,0 +1,162 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Tests for python.compiler.mlir."""
+
+from tensorflow.python.compiler.mlir import mlir
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.pywrap_mlir import import_graphdef
+
+
+class MLIRGraphDefImportTest(test.TestCase):
+
+  def testImport(self):
+    """Tests the basic flow of `tf.mlir.experimental.convert_graph_def`."""
+    mlir_module = mlir.convert_graph_def('')
+    # An empty graph should contain at least an empty main function.
+    self.assertIn('func @main', mlir_module)
+
+  def testInvalidPbtxt(self):
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not parse input proto'):
+      mlir.convert_graph_def('some invalid proto')
+
+  def testGraphDefToTf(self):
+    """Tests the basic flow of `tf.mlir.experimental.convert_graph_def`
+
+        with tf-standard-pipeline converting all the way to the TF dialect.
+    """
+
+    tensor_shape = (10, 10)
+
+    @def_function.function(
+        input_signature=(
+            tensor_spec.TensorSpec(shape=tensor_shape, dtype=dtypes.float32),
+            tensor_spec.TensorSpec(shape=tensor_shape, dtype=dtypes.float32),
+        ))
+    def add_func(lhs, rhs):
+      return math_ops.add(lhs, rhs)
+
+    tf_graph_def = add_func.get_concrete_function().graph.as_graph_def()
+
+    mlir_tf = import_graphdef(
+        tf_graph_def,
+        "tf-standard-pipeline",
+        False,
+        input_names=["lhs", "rhs"],
+        input_data_types=["DT_FLOAT", "DT_FLOAT"],
+        input_data_shapes=["10,10", "10,10"],
+        output_names=["Add"])
+    # Check whether the mlir-function signature has the mentioned
+    # inputs and outputs.
+    self.assertRegex(
+        mlir_tf,
+        r"func @main\(%arg0: tensor<10x10xf32>, %arg1: tensor<10x10xf32>")
+    self.assertRegex(mlir_tf, r'inputs = "lhs,rhs"')
+    self.assertRegex(mlir_tf, r'outputs = "Add"')
+
+    # Same check with scalar input (empty input shape).
+    mlir_tf = import_graphdef(
+        tf_graph_def,
+        "tf-standard-pipeline",
+        False,
+        input_names=["lhs", "rhs"],
+        input_data_types=["DT_FLOAT", "DT_FLOAT"],
+        input_data_shapes=["", ""],
+        output_names=["Add"])
+    self.assertRegex(mlir_tf,
+                     r"func @main\(%arg0: tensor<f32>, %arg1: tensor<f32>")
+
+    # Test invalid test cases where no. of input names is invalid/wrong.
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Length of input node array and data type doesn't match"):
+
+      import_graphdef(
+          tf_graph_def,
+          "tf-standard-pipeline",
+          False,
+          input_names=["lhs"],
+          input_data_types=["DT_FLOAT", "DT_FLOAT"],
+          input_data_shapes=["10,10", "10,10"],
+          output_names=["Add"])
+
+    # Test invalid test cases where the input shapes argument is wrong.
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Dimensions must be equal"):
+
+      import_graphdef(
+          tf_graph_def,
+          "tf-standard-pipeline",
+          False,
+          input_names=["lhs", "rhs"],
+          input_data_types=["DT_FLOAT", "DT_FLOAT"],
+          input_data_shapes=["10,11", "10,10"],
+          output_names=["Add"])
+
+
+class MLIRConcreteFunctionImportTest(test.TestCase):
+
+  @test_util.run_v2_only
+  def testImport(self):
+
+    @def_function.function
+    def sqr(i):
+      return i * i
+
+    concrete_function = sqr.get_concrete_function(
+        tensor_spec.TensorSpec(None, dtypes.float32))
+    mlir_module = mlir.convert_function(concrete_function, show_debug_info=True)
+    self.assertRegex(mlir_module, r'func @.*sqr.*\(')
+    self.assertRegex(mlir_module, r'loc\(".*mlir_test.py":.*:1\)')
+
+  @test_util.run_v2_only
+  def testImportWithCall(self):
+
+    @def_function.function
+    def callee(i):
+      return i
+
+    @def_function.function
+    def caller(i):
+      return callee(i)
+
+    concrete_function = caller.get_concrete_function(
+        tensor_spec.TensorSpec(None, dtypes.float32))
+    mlir_module = mlir.convert_function(concrete_function)
+    self.assertRegex(mlir_module, r'func @.*caller.*\(')
+    self.assertRegex(mlir_module, r'func private @.*callee.*\(')
+
+  @test_util.run_v2_only
+  def testImportWithControlRet(self):
+
+    @def_function.function
+    def logging():
+      logging_ops.print_v2('some message')
+
+    concrete_function = logging.get_concrete_function()
+    mlir_module = mlir.convert_function(concrete_function, pass_pipeline='')
+    self.assertRegex(mlir_module, r'tf\.PrintV2')
+    self.assertRegex(mlir_module, r'tf_executor.fetch.*: !tf_executor.control')
+
+
+if __name__ == '__main__':
+  test.main()