From cf3c40795dfd3749d154b0143f90087a624533e3 Mon Sep 17 00:00:00 2001 From: Mehdi Amini Date: Mon, 9 Sep 2019 14:08:10 -0700 Subject: [PATCH] Implement a basic API: `tf.mlir.experimental.convert_graph_def` Load a .pbptx, convert to MLIR, and (optionally) optimize the module before returning it as a string. This is an early experimental API, intended for example to play with some some colab examples during development. PiperOrigin-RevId: 268073235 --- tensorflow/compiler/mlir/python/BUILD | 11 +++ tensorflow/compiler/mlir/python/mlir.i | 74 +++++++++++++++++++ .../mlir/tensorflow/translate/import_model.cc | 9 +++ .../mlir/tensorflow/translate/import_model.h | 5 ++ .../mlir/tensorflow/utils/import_utils.cc | 46 +++++++----- .../mlir/tensorflow/utils/import_utils.h | 5 ++ tensorflow/python/BUILD | 1 + tensorflow/python/__init__.py | 3 + tensorflow/python/compiler/BUILD | 1 + tensorflow/python/compiler/mlir/BUILD | 27 +++++++ tensorflow/python/compiler/mlir/mlir.py | 38 ++++++++++ tensorflow/python/compiler/mlir/mlir_test.py | 41 ++++++++++ tensorflow/python/tensorflow.i | 2 + .../tools/api/generator/api_init_files.bzl | 2 + .../tools/api/generator/api_init_files_v1.bzl | 2 + .../v1/tensorflow.mlir.experimental.pbtxt | 7 ++ .../tools/api/golden/v1/tensorflow.mlir.pbtxt | 7 ++ .../tools/api/golden/v1/tensorflow.pbtxt | 4 + .../v2/tensorflow.mlir.experimental.pbtxt | 7 ++ .../tools/api/golden/v2/tensorflow.mlir.pbtxt | 7 ++ .../tools/api/golden/v2/tensorflow.pbtxt | 4 + 21 files changed, 284 insertions(+), 19 deletions(-) create mode 100644 tensorflow/compiler/mlir/python/BUILD create mode 100644 tensorflow/compiler/mlir/python/mlir.i create mode 100644 tensorflow/python/compiler/mlir/BUILD create mode 100644 tensorflow/python/compiler/mlir/mlir.py create mode 100644 tensorflow/python/compiler/mlir/mlir_test.py create mode 100644 tensorflow/tools/api/golden/v1/tensorflow.mlir.experimental.pbtxt create mode 100644 tensorflow/tools/api/golden/v1/tensorflow.mlir.pbtxt create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.mlir.experimental.pbtxt create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.mlir.pbtxt diff --git a/tensorflow/compiler/mlir/python/BUILD b/tensorflow/compiler/mlir/python/BUILD new file mode 100644 index 00000000000..5291cf3b141 --- /dev/null +++ b/tensorflow/compiler/mlir/python/BUILD @@ -0,0 +1,11 @@ +package( + default_visibility = ["//visibility:public"], + licenses = ["notice"], # Apache 2.0 +) + +exports_files( + ["mlir.i"], + visibility = [ + "//tensorflow/python:__subpackages__", + ], +) diff --git a/tensorflow/compiler/mlir/python/mlir.i b/tensorflow/compiler/mlir/python/mlir.i new file mode 100644 index 00000000000..03273357b2b --- /dev/null +++ b/tensorflow/compiler/mlir/python/mlir.i @@ -0,0 +1,74 @@ +/* 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. +==============================================================================*/ + +%include "tensorflow/python/platform/base.i" + +%{ + +#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h" +#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h" +#include "tensorflow/compiler/mlir/tensorflow/utils/import_utils.h" + +namespace tensorflow { +namespace swig { + +// Simple wrapper to support tf.mlir.experimental.convert_graph_def. +// Load a .pbptx, convert to MLIR, and (optionally) optimize the module before +// returning it as a string. +// This is an early experimental API, ideally we should return a wrapper object +// around a Python binding to the MLIR module. +string ImportGraphDef(const string &proto, TF_Status* status) { + GraphDef graphdef; + auto s = tensorflow::LoadProtoFromBuffer(proto, &graphdef); + if (!s.ok()) { + Set_TF_Status_from_Status(status, s); + return "// error"; + } + GraphDebugInfo debug_info; + NodeSpecs specs; + mlir::MLIRContext context; + auto module = ConvertGraphdefToMlir(graphdef, debug_info, specs, &context); + if (!module.ok()) { + Set_TF_Status_from_Status(status, module.status()); + return "// error"; + } + + return MlirModuleToString(*module.ConsumeValueOrDie()); +} + +} // namespace swig +} // namespace tensorflow + +%} + +%ignoreall + +%unignore tensorflow; +%unignore tensorflow::swig; +%unignore tensorflow::swig::ImportGraphDef; + +// Wrap this function +namespace tensorflow { +namespace swig { +static string ImportGraphDef(const string &graphdef, TF_Status* status); +} // namespace swig +} // namespace tensorflow + +%insert("python") %{ +def import_graphdef(graphdef): + return str(ImportGraphDef(str(graphdef).encode('utf-8'))); +%} + +%unignoreall diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc index e7bfd578046..c06bd3ec5c3 100644 --- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc +++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc @@ -1774,4 +1774,13 @@ StatusOr ConvertSavedModelToMlir( add_default_attributes, context); } +std::string MlirModuleToString(mlir::ModuleOp module) { + std::string txt_module; + { + llvm::raw_string_ostream os{txt_module}; + module.print(os); + } + return txt_module; +} + } // namespace tensorflow diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.h b/tensorflow/compiler/mlir/tensorflow/translate/import_model.h index 98bb607fa6a..6ca4c0098d7 100644 --- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.h +++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.h @@ -16,6 +16,8 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_IMPORT_MODEL_H_ #define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_IMPORT_MODEL_H_ +#include + #include "mlir/IR/MLIRContext.h" // TF:local_config_mlir #include "mlir/IR/Module.h" // TF:local_config_mlir #include "tensorflow/cc/saved_model/loader.h" @@ -48,6 +50,9 @@ stream_executor::port::StatusOr ConvertSavedModelToMlir( const SavedModelBundle& saved_model, const GraphDebugInfo& debug_info, mlir::MLIRContext* context, bool add_default_attributes = true); +// Serialize a MLIR module to a string. +std::string MlirModuleToString(mlir::ModuleOp m); + } // namespace tensorflow #endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_IMPORT_MODEL_H_ diff --git a/tensorflow/compiler/mlir/tensorflow/utils/import_utils.cc b/tensorflow/compiler/mlir/tensorflow/utils/import_utils.cc index 52fb7cac5b7..5be0ebd6894 100644 --- a/tensorflow/compiler/mlir/tensorflow/utils/import_utils.cc +++ b/tensorflow/compiler/mlir/tensorflow/utils/import_utils.cc @@ -37,6 +37,30 @@ inline llvm::StringRef StringViewToRef(absl::string_view view) { namespace tensorflow { +Status LoadProtoFromBuffer(absl::string_view input, + tensorflow::protobuf::Message* proto) { + tensorflow::protobuf::TextFormat::Parser parser; + // Don't produce errors when attempting to parse text format as it would fail + // when the input is actually a binary file. + NoOpErrorCollector collector; + parser.RecordErrorsTo(&collector); + // Attempt to parse as text. + tensorflow::protobuf::io::ArrayInputStream input_stream(input.data(), + input.size()); + if (parser.Parse(&input_stream, proto)) { + return Status::OK(); + } + // Else attempt to parse as binary. + proto->Clear(); + tensorflow::protobuf::io::ArrayInputStream binary_stream(input.data(), + input.size()); + if (proto->ParseFromZeroCopyStream(&binary_stream)) { + return Status::OK(); + } + LOG(ERROR) << "Error parsing Protobuf"; + return errors::InvalidArgument("Could not parse input proto"); +} + Status LoadProtoFromFile(absl::string_view input_filename, tensorflow::protobuf::Message* proto) { auto file_or_err = @@ -45,26 +69,10 @@ Status LoadProtoFromFile(absl::string_view input_filename, return errors::InvalidArgument("Could not open input file"); auto& input_file = *file_or_err; - std::string content(input_file->getBufferStart(), - input_file->getBufferSize()); + absl::string_view content(input_file->getBufferStart(), + input_file->getBufferSize()); - tensorflow::protobuf::TextFormat::Parser parser; - // Don't produce errors when attempting to parse text format as it would fail - // when the input is actually a binary file. - NoOpErrorCollector collector; - parser.RecordErrorsTo(&collector); - // Attempt to parse as text. - if (parser.ParseFromString(content, proto)) { - return Status::OK(); - } - // Else attempt to parse as binary. - proto->Clear(); - std::istringstream istream(content); - if (proto->ParseFromIstream(&istream)) { - return Status::OK(); - } - LOG(ERROR) << "Error parsing Protobuf: " << input_filename; - return errors::InvalidArgument("Could not parse input file"); + return LoadProtoFromBuffer(content, proto); } } // namespace tensorflow diff --git a/tensorflow/compiler/mlir/tensorflow/utils/import_utils.h b/tensorflow/compiler/mlir/tensorflow/utils/import_utils.h index 1158b9a6173..a7d00cf890e 100644 --- a/tensorflow/compiler/mlir/tensorflow/utils/import_utils.h +++ b/tensorflow/compiler/mlir/tensorflow/utils/import_utils.h @@ -22,6 +22,11 @@ limitations under the License. namespace tensorflow { +// Reads text (.pbtext) or binary (.pb) format of a proto message from the given +// buffer. Returns error status of the file is not found or malformed proto. +Status LoadProtoFromBuffer(absl::string_view input, + tensorflow::protobuf::Message* proto); + // Reads text (.pbtext) or binary (.pb) format of a proto message from the given // file path. Returns error status of the file is not found or malformed proto. Status LoadProtoFromFile(absl::string_view input_filename, diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index db2173c3c06..a94bc6dcdc9 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -5029,6 +5029,7 @@ tf_py_wrap_cc( "util/scoped_annotation.i", "util/traceme.i", "util/transform_graph.i", + "//tensorflow/compiler/mlir/python:mlir.i", "//tensorflow/lite/toco/python:toco.i", ], # add win_def_file for pywrap_tensorflow diff --git a/tensorflow/python/__init__.py b/tensorflow/python/__init__.py index d733dc0d93b..e06d7a9af39 100644 --- a/tensorflow/python/__init__.py +++ b/tensorflow/python/__init__.py @@ -160,6 +160,9 @@ from tensorflow.python.ops import rnn_cell from tensorflow.python.compiler.xla import jit from tensorflow.python.compiler.xla import xla +# MLIR APIs. +from tensorflow.python.compiler.mlir import mlir + # Required due to `rnn` and `rnn_cell` not being imported in `nn` directly # (due to a circular dependency issue: rnn depends on layers). nn.dynamic_rnn = rnn.dynamic_rnn diff --git a/tensorflow/python/compiler/BUILD b/tensorflow/python/compiler/BUILD index 9d3aa19d44d..4d16a85e379 100644 --- a/tensorflow/python/compiler/BUILD +++ b/tensorflow/python/compiler/BUILD @@ -17,6 +17,7 @@ py_library( deps = if_not_windows([ "//tensorflow/python/compiler/tensorrt:init_py", ]) + [ + "//tensorflow/python/compiler/mlir", "//tensorflow/python/compiler/xla:compiler_py", ], ) diff --git a/tensorflow/python/compiler/mlir/BUILD b/tensorflow/python/compiler/mlir/BUILD new file mode 100644 index 00000000000..ee191e4e6db --- /dev/null +++ b/tensorflow/python/compiler/mlir/BUILD @@ -0,0 +1,27 @@ +load("//tensorflow:tensorflow.bzl", "py_test") + +package( + default_visibility = ["//visibility:public"], + licenses = ["notice"], # Apache 2.0 +) + +py_library( + name = "mlir", + srcs = ["mlir.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:pywrap_tensorflow", + "//tensorflow/python:util", + ], +) + +py_test( + name = "mlir_test", + srcs = ["mlir_test.py"], + python_version = "PY3", + deps = [ + ":mlir", + "//tensorflow/python:client_testlib", + "//tensorflow/python:platform", + ], +) diff --git a/tensorflow/python/compiler/mlir/mlir.py b/tensorflow/python/compiler/mlir/mlir.py new file mode 100644 index 00000000000..42a9971e06a --- /dev/null +++ b/tensorflow/python/compiler/mlir/mlir.py @@ -0,0 +1,38 @@ +# 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 __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python import pywrap_tensorflow as import_graphdef +from tensorflow.python.util.tf_export import tf_export + + +@tf_export('mlir.experimental.convert_graph_def') +def convert_graph_def(graph_def): + """Import a GraphDef and convert it to a textual MLIR module. + + Args: + graph_def: An object of type graph_pb2.GraphDef or a textual proto + representation of a valid GraphDef. + + Returns: + A textual representation of the MLIR module corresponding to the graphdef. + Raises a RuntimeError on error. + + """ + return import_graphdef.import_graphdef(graph_def) diff --git a/tensorflow/python/compiler/mlir/mlir_test.py b/tensorflow/python/compiler/mlir/mlir_test.py new file mode 100644 index 00000000000..8fd5952fd76 --- /dev/null +++ b/tensorflow/python/compiler/mlir/mlir_test.py @@ -0,0 +1,41 @@ +# 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 __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.compiler.mlir import mlir +from tensorflow.python.framework import errors +from tensorflow.python.platform import test + + +class MLIRImportTest(test.TestCase): + + def test_import_graph_def(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 test_invalid_pbtxt(self): + with self.assertRaisesRegexp(errors.InvalidArgumentError, + 'Could not parse input proto'): + mlir.convert_graph_def('some invalid proto') + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/python/tensorflow.i b/tensorflow/python/tensorflow.i index 9ab816978f8..6c3acd6b293 100644 --- a/tensorflow/python/tensorflow.i +++ b/tensorflow/python/tensorflow.i @@ -55,3 +55,5 @@ limitations under the License. %include "tensorflow/python/util/traceme.i" %include "tensorflow/python/util/scoped_annotation.i" + +%include "tensorflow/compiler/mlir/python/mlir.i" diff --git a/tensorflow/python/tools/api/generator/api_init_files.bzl b/tensorflow/python/tools/api/generator/api_init_files.bzl index 38dfff6525f..32a6554518d 100644 --- a/tensorflow/python/tools/api/generator/api_init_files.bzl +++ b/tensorflow/python/tools/api/generator/api_init_files.bzl @@ -39,6 +39,8 @@ TENSORFLOW_API_INIT_FILES = [ "math/__init__.py", "mixed_precision/__init__.py", "mixed_precision/experimental/__init__.py", + "mlir/__init__.py", + "mlir/experimental/__init__.py", "nest/__init__.py", "nn/__init__.py", "quantization/__init__.py", diff --git a/tensorflow/python/tools/api/generator/api_init_files_v1.bzl b/tensorflow/python/tools/api/generator/api_init_files_v1.bzl index 94d72c2a878..dfd1f12c8f2 100644 --- a/tensorflow/python/tools/api/generator/api_init_files_v1.bzl +++ b/tensorflow/python/tools/api/generator/api_init_files_v1.bzl @@ -49,6 +49,8 @@ TENSORFLOW_API_INIT_FILES_V1 = [ "manip/__init__.py", "math/__init__.py", "metrics/__init__.py", + "mlir/__init__.py", + "mlir/experimental/__init__.py", "nest/__init__.py", "nn/__init__.py", "nn/rnn_cell/__init__.py", diff --git a/tensorflow/tools/api/golden/v1/tensorflow.mlir.experimental.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.mlir.experimental.pbtxt new file mode 100644 index 00000000000..b39c6d3b85d --- /dev/null +++ b/tensorflow/tools/api/golden/v1/tensorflow.mlir.experimental.pbtxt @@ -0,0 +1,7 @@ +path: "tensorflow.mlir.experimental" +tf_module { + member_method { + name: "convert_graph_def" + argspec: "args=[\'graph_def\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v1/tensorflow.mlir.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.mlir.pbtxt new file mode 100644 index 00000000000..d4ad55c8999 --- /dev/null +++ b/tensorflow/tools/api/golden/v1/tensorflow.mlir.pbtxt @@ -0,0 +1,7 @@ +path: "tensorflow.mlir" +tf_module { + member { + name: "experimental" + mtype: "" + } +} diff --git a/tensorflow/tools/api/golden/v1/tensorflow.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.pbtxt index aa742a5e1c9..976da90a72d 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.pbtxt @@ -508,6 +508,10 @@ tf_module { name: "metrics" mtype: "" } + member { + name: "mlir" + mtype: "" + } member { name: "name_scope" mtype: "" diff --git a/tensorflow/tools/api/golden/v2/tensorflow.mlir.experimental.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.mlir.experimental.pbtxt new file mode 100644 index 00000000000..b39c6d3b85d --- /dev/null +++ b/tensorflow/tools/api/golden/v2/tensorflow.mlir.experimental.pbtxt @@ -0,0 +1,7 @@ +path: "tensorflow.mlir.experimental" +tf_module { + member_method { + name: "convert_graph_def" + argspec: "args=[\'graph_def\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v2/tensorflow.mlir.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.mlir.pbtxt new file mode 100644 index 00000000000..d4ad55c8999 --- /dev/null +++ b/tensorflow/tools/api/golden/v2/tensorflow.mlir.pbtxt @@ -0,0 +1,7 @@ +path: "tensorflow.mlir" +tf_module { + member { + name: "experimental" + mtype: "" + } +} diff --git a/tensorflow/tools/api/golden/v2/tensorflow.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.pbtxt index ee3c0cc22bb..fec187d73cd 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.pbtxt @@ -260,6 +260,10 @@ tf_module { name: "mixed_precision" mtype: "" } + member { + name: "mlir" + mtype: "" + } member { name: "name_scope" mtype: ""