Add XLA compiler frontend. * Names it iree-import-xla (backwards from iree-tf-import) because Ben asked me to swap the order. I'll followup with a rename of the TF side. * Does not test HLO pbtext because I couldn't find a reliable way to generate it from OSS tooling. * Fixes a little bug in the bazel integration that was keeping more than one target from building.
diff --git a/CMakeLists.txt b/CMakeLists.txt index 9add1e0..0aaf142 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt
@@ -45,15 +45,24 @@ option(IREE_BUILD_DEBUGGER "Builds the IREE debugger app." OFF) option(IREE_BUILD_JAVA_BINDINGS "Builds the IREE java bindings." OFF) option(IREE_BUILD_EXPERIMENTAL "Builds experimental projects." OFF) -option(IREE_BUILD_TENSORFLOW_COMPILER "Builds TensorFlow compiler." OFF) +option(IREE_BUILD_TENSORFLOW_COMPILER "Builds TensorFlow compiler frontend." OFF) +option(IREE_BUILD_XLA_COMPILER "Builds TensorFlow XLA compiler frontend." OFF) set(IREE_HAL_DRIVERS_TO_BUILD "all" CACHE STRING "Semicolon-separated list of HAL drivers to build, or \"all\".") set(IREE_TARGET_BACKENDS_TO_BUILD "all" CACHE STRING "Semicolon-separated list of target backends to build, or \"all\".") +# Master enable for tensorflow build support. +# Note that this is a normal CMake variable used to gate build features (not +# a cache variable that is user-settable). +set(IREE_ENABLE_TENSORFLOW OFF) +if(${IREE_BUILD_TENSORFLOW_COMPILER} OR ${IREE_BUILD_XLA_COMPILER}) + set(IREE_ENABLE_TENSORFLOW ON) +endif() + # Default python bindings to enabled for some features. -if(${IREE_BUILD_TENSORFLOW_COMPILER}) +if(${IREE_ENABLE_TENSORFLOW}) option(IREE_BUILD_PYTHON_BINDINGS "Builds the IREE python bindings" ON) else() option(IREE_BUILD_PYTHON_BINDINGS "Builds the IREE python bindings" OFF) @@ -282,8 +291,7 @@ #------------------------------------------------------------------------------- if(${IREE_BUILD_COMPILER} OR - ${IREE_BUILD_PYTHON_BINDINGS} OR - ${IREE_BUILD_TENSORFLOW_COMPILER}) + ${IREE_BUILD_PYTHON_BINDINGS}) find_package(Python3 COMPONENTS Interpreter REQUIRED) endif() @@ -404,7 +412,7 @@ # Depends on python configuration. #------------------------------------------------------------------------------- -if(${IREE_BUILD_TENSORFLOW_COMPILER}) +if(${IREE_ENABLE_TENSORFLOW}) include(configure_bazel) iree_configure_bazel() endif() @@ -531,7 +539,7 @@ add_subdirectory(experimental) endif() -if(${IREE_BUILD_TENSORFLOW_COMPILER}) +if(${IREE_ENABLE_TENSORFLOW}) add_subdirectory(integrations/tensorflow) endif()
diff --git a/bindings/python/pyiree/compiler2/CMakeLists.txt b/bindings/python/pyiree/compiler2/CMakeLists.txt index c9a4ef8..f5bf412 100644 --- a/bindings/python/pyiree/compiler2/CMakeLists.txt +++ b/bindings/python/pyiree/compiler2/CMakeLists.txt
@@ -23,4 +23,5 @@ "core.py" "tf.py" "tools.py" + "xla.py" )
diff --git a/bindings/python/pyiree/compiler2/tools.py b/bindings/python/pyiree/compiler2/tools.py index 1d32a15..76df5dc 100644 --- a/bindings/python/pyiree/compiler2/tools.py +++ b/bindings/python/pyiree/compiler2/tools.py
@@ -37,6 +37,7 @@ # a python module that provides a `get_tool` function for getting its absolute # path. This dictionary maps the tool name to the module. _TOOL_MODULE_MAP = { + "iree-import-xla": "pyiree.tools.xla", "iree-tf-import": "pyiree.tools.tf", "iree-translate": "pyiree.tools.core", } @@ -46,6 +47,7 @@ _TOOL_MODULE_PACKAGES = { "pyiree.tools.core": "google-iree-tools-core", "pyiree.tools.tf": "google-iree-tools-tf", + "pyiree.tools.xla": "google-iree-tools-xla", } # Environment variable holding directories to be searched for named tools. @@ -163,7 +165,7 @@ input_file_handle.close() -def invoke_pipeline(command_lines: List[List[str]]): +def invoke_pipeline(command_lines: List[List[str]], immediate_input=None): """Invoke a pipeline of commands. The first stage of the pipeline will have its stdin set to DEVNULL and each @@ -173,7 +175,9 @@ an exception raised with its stderr output. """ stages = [] - prev_out = subprocess.DEVNULL + pipeline_input = (subprocess.DEVNULL + if immediate_input is None else subprocess.PIPE) + prev_out = pipeline_input stderr_handle = sys.stderr # Create all stages. @@ -193,6 +197,16 @@ for stage in stages: stage.start() + # Pump input. + pipe_success = True + if immediate_input is not None: + try: + pipe_success = False + stages[0].process.stdin.write(immediate_input) + pipe_success = True + finally: + stages[0].process.stdin.close() + # Join. for stage in stages: stage.join() @@ -203,6 +217,10 @@ if stage.completed.returncode != 0: raise CompilerToolError(stage.completed) + # Broken pipe. + if not pipe_success: + raise CompilerToolError(stages[0].completed) + # Print any stderr output. for stage in stages: _write_binary_stderr(stderr_handle, stage.errs)
diff --git a/bindings/python/pyiree/compiler2/xla.py b/bindings/python/pyiree/compiler2/xla.py new file mode 100644 index 0000000..4f93937 --- /dev/null +++ b/bindings/python/pyiree/compiler2/xla.py
@@ -0,0 +1,184 @@ +# Lint-as: python3 +"""XLA compiler interface.""" + +# Copyright 2020 Google LLC +# +# 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 +# +# https://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. + +# TODO(#4131) python>=3.7: Use postponed type annotations. + +from enum import Enum +import logging +import tempfile +from typing import List, Optional, Sequence, Set, Union + +from .tools import find_tool, invoke_immediate, invoke_pipeline +from .core import CompilerOptions, DEFAULT_TESTING_BACKENDS, build_compile_command_line + +__all__ = [ + "compile_file", + "compile_str", + "is_available", + "DEFAULT_TESTING_BACKENDS", + "ImportOptions", + "ImportFormat", +] + +_IMPORT_TOOL = "iree-import-xla" + + +def is_available(): + """Determine if the XLA frontend is available.""" + try: + find_tool(_IMPORT_TOOL) + except ValueError: + logging.warning("Unable to find IREE tool %s", _IMPORT_TOOL) + return False + return True + + +class ImportFormat(Enum): + """Import type of the model.""" + BINARY_PROTO = "binary_proto" + TEXT_PROTO = "text_proto" + + @staticmethod + def parse(spec: Union[str, "ImportFormat"]) -> "ImportFormat": + """Parses or returns an ImportFormat. + + Args: + spec: An ImportFormat instance or the case-insensitive name of one of + the enum values. + Returns: + An ImportFormat instance. + """ + if isinstance(spec, ImportFormat): + return spec + spec = spec.upper() + if spec not in ImportFormat.__members__: + raise ValueError(f"For import_format= argument, expected one of: " + f"{', '.join(ImportFormat.__members__.keys())}") + return ImportFormat[spec] + + +# TODO(#4131) python>=3.7: Consider using a dataclass. +class ImportOptions(CompilerOptions): + """Import options layer on top of the backend compiler options.""" + + def __init__(self, + import_only: bool = False, + import_format: Union[ImportFormat, + str] = ImportFormat.BINARY_PROTO, + import_extra_args: Sequence[str] = (), + save_temp_iree_input: Optional[str] = None, + **kwargs): + """Initialize options from keywords. + + Args: + import_format: Format of the proto (text or binary). + save_temp_iree_input: Optionally save the IR that is the result of the + import (ready to be passed to IREE). + """ + super().__init__(**kwargs) + self.import_only = import_only + self.import_format = ImportFormat.parse(import_format) + self.import_extra_args = import_extra_args + self.save_temp_iree_input = save_temp_iree_input + + +def build_import_command_line(input_path: str, + options: ImportOptions) -> List[str]: + """Builds a command line for invoking the import stage. + + Args: + input_path: The input path. + options: Import options. + Returns: + List of strings of command line. + """ + import_tool = find_tool(_IMPORT_TOOL) + cl = [ + import_tool, + input_path, + f"--xla-format={options.import_format.value}", + ] + if options.import_only and options.output_file: + # Import stage directly outputs. + if options.output_file: + cl.append(f"-o={options.output_file}") + # Save temps flags. + if options.save_temp_iree_input: + cl.append(f"--save-temp-iree-input={options.save_temp_iree_input}") + # Extra args. + cl.extend(options.import_extra_args) + return cl + + +def compile_file(xla_file_path: str, **kwargs): + """Compiles an on-disk XLA protocol buffer to an IREE binary. + + Args: + xla_file_path: Path to the XLA proto file. + **kwargs: Keyword args corresponding to ImportOptions or CompilerOptions. + Returns: + A bytes-like object with the compiled output or None if output_file= + was specified. + """ + options = ImportOptions(**kwargs) + import_cl = build_import_command_line(xla_file_path, options) + if options.import_only: + # One stage tool pipeline. + result = invoke_immediate(import_cl) + if options.output_file: + return None + return result + + # Full compilation pipeline. + compile_cl = build_compile_command_line("-", options) + result = invoke_pipeline([import_cl, compile_cl]) + if options.output_file: + return None + return result + + +def compile_str(xla_content: Union[bytes, str], **kwargs): + """Compiles in-memory XLA content to an IREE binary. + + Args: + xla_content: Either bytes or str content (str is only valid for text + formats). + **kwargs: Keyword args corresponding to ImportOptions or CompilerOptions. + Returns: + A bytes-like object with the compiled output or None if output_file= + was specified. + """ + options = ImportOptions(**kwargs) + if isinstance(xla_content, str): + if options.import_format != ImportFormat.TEXT_PROTO: + raise ValueError("If passing a string, ImportFormat must be TEXT_PROTO") + xla_content = xla_content.encode("utf-8") + + import_cl = build_import_command_line("-", options) + if options.import_only: + # One stage tool pipeline. + result = invoke_immediate(import_cl, immediate_input=xla_content) + if options.output_file: + return None + return result + + # Full compilation pipeline. + compile_cl = build_compile_command_line("-", options) + result = invoke_pipeline([import_cl, compile_cl], immediate_input=xla_content) + if options.output_file: + return None + return result
diff --git a/bindings/python/tests/CMakeLists.txt b/bindings/python/tests/CMakeLists.txt index 701bb72..2e1f330 100644 --- a/bindings/python/tests/CMakeLists.txt +++ b/bindings/python/tests/CMakeLists.txt
@@ -25,3 +25,10 @@ SRCS "compiler_tf_test.py" ) + +iree_py_test( + NAME + compiler_xla_test + SRCS + "compiler_xla_test.py" +)
diff --git a/bindings/python/tests/compiler_xla_test.py b/bindings/python/tests/compiler_xla_test.py new file mode 100644 index 0000000..cd8f241 --- /dev/null +++ b/bindings/python/tests/compiler_xla_test.py
@@ -0,0 +1,95 @@ +# Lint as: python3 +# Copyright 2020 Google LLC +# +# 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 +# +# https://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. + +import logging +import os +import sys +import tempfile +import unittest + +# TODO: No idea why pytype cannot find names from this module. +# pytype: disable=name-error +from pyiree.compiler2.xla import * + +if not is_available(): + print(f"Skipping test {__file__} because the IREE XLA compiler " + f"is not installed") + sys.exit(0) + + +class CompilerTest(unittest.TestCase): + + def testImportBinaryPbFile(self): + path = os.path.join(os.path.dirname(__file__), "testdata", "xla_sample.pb") + text = compile_file(path, import_only=True).decode("utf-8") + logging.info("%s", text) + self.assertIn("mhlo.constant", text) + + def testCompileBinaryPbFile(self): + path = os.path.join(os.path.dirname(__file__), "testdata", "xla_sample.pb") + binary = compile_file(path, target_backends=DEFAULT_TESTING_BACKENDS) + logging.info("Binary length = %d", len(binary)) + self.assertIn(b"main", binary) + + def testImportBinaryPbFileOutputFile(self): + path = os.path.join(os.path.dirname(__file__), "testdata", "xla_sample.pb") + with tempfile.NamedTemporaryFile("wt", delete=False) as f: + try: + f.close() + output = compile_file(path, import_only=True, output_file=f.name) + self.assertIsNone(output) + with open(f.name, "rt") as f_read: + text = f_read.read() + finally: + os.remove(f.name) + logging.info("%s", text) + self.assertIn("mhlo.constant", text) + + def testCompileBinaryPbFileOutputFile(self): + path = os.path.join(os.path.dirname(__file__), "testdata", "xla_sample.pb") + with tempfile.NamedTemporaryFile("wt", delete=False) as f: + try: + f.close() + output = compile_file(path, + output_file=f.name, + target_backends=DEFAULT_TESTING_BACKENDS) + self.assertIsNone(output) + with open(f.name, "rb") as f_read: + binary = f_read.read() + finally: + os.remove(f.name) + logging.info("Binary length = %d", len(binary)) + self.assertIn(b"main", binary) + + def testImportBinaryPbBytes(self): + path = os.path.join(os.path.dirname(__file__), "testdata", "xla_sample.pb") + with open(path, "rb") as f: + content = f.read() + text = compile_str(content, import_only=True).decode("utf-8") + logging.info("%s", text) + self.assertIn("mhlo.constant", text) + + def testCompileBinaryPbBytes(self): + path = os.path.join(os.path.dirname(__file__), "testdata", "xla_sample.pb") + with open(path, "rb") as f: + content = f.read() + binary = compile_str(content, target_backends=DEFAULT_TESTING_BACKENDS) + logging.info("Binary length = %d", len(binary)) + self.assertIn(b"main", binary) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.DEBUG) + unittest.main()
diff --git a/bindings/python/tests/testdata/generate_xla.py b/bindings/python/tests/testdata/generate_xla.py new file mode 100644 index 0000000..00f872f --- /dev/null +++ b/bindings/python/tests/testdata/generate_xla.py
@@ -0,0 +1,34 @@ +# Copyright 2020 Google LLC +# +# 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 +# +# https://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. + +import os + +import numpy as np + +# Jax is the most accessible way to get at an xla_client. +# python -m pip install --upgrade pip +# python -m pip install --upgrade jax jaxlib +from jaxlib import xla_client + +ops = xla_client.ops + +builder = xla_client.XlaBuilder("testbuilder") +in_shape = np.array([4], dtype=np.float32) +in_feed = ops.Parameter(builder, 0, xla_client.shape_from_pyval(in_shape)) +result = ops.Add(in_feed, ops.Constant(builder, np.float32(1.0))) +xla_computation = builder.Build(result) + +this_dir = os.path.dirname(__file__) +with open(os.path.join(this_dir, "xla_sample.pb"), "wb") as f: + f.write(xla_computation.as_serialized_hlo_module_proto())
diff --git a/bindings/python/tests/testdata/xla_sample.pb b/bindings/python/tests/testdata/xla_sample.pb new file mode 100644 index 0000000..1f69a64 --- /dev/null +++ b/bindings/python/tests/testdata/xla_sample.pb Binary files differ
diff --git a/build_tools/cmake/configure_bazel.cmake b/build_tools/cmake/configure_bazel.cmake index 528d244..6667a9d 100644 --- a/build_tools/cmake/configure_bazel.cmake +++ b/build_tools/cmake/configure_bazel.cmake
@@ -137,7 +137,7 @@ add_custom_target(${ARG_INVOCATION_TARGET} USES_TERMINAL COMMAND ${CMAKE_COMMAND} -E echo - "Starting bazel build of targets ${ARG_BAZEL_TARGETS}" + "Starting bazel build of targets '${ARG_BAZEL_TARGETS}'" COMMAND "${IREE_BAZEL_WRAPPER}" build ${ARG_BAZEL_TARGETS} COMMAND ${CMAKE_COMMAND} -E echo "Bazel build complete." )
diff --git a/integrations/tensorflow/CMakeLists.txt b/integrations/tensorflow/CMakeLists.txt index ba528f1..f54fd26 100644 --- a/integrations/tensorflow/CMakeLists.txt +++ b/integrations/tensorflow/CMakeLists.txt
@@ -22,13 +22,24 @@ # If this directory is included, then building TensorFlow is assumed (the # config option happens at the higher level). +set(_bazel_targets) +set(_executable_paths) + +if(${IREE_BUILD_TENSORFLOW_COMPILER}) + list(APPEND _bazel_targets //integrations/tensorflow/compiler:iree-tf-import) + list(APPEND _executable_paths integrations/tensorflow/compiler/iree-tf-import) +endif() + +if(${IREE_BUILD_XLA_COMPILER}) + list(APPEND _bazel_targets //integrations/tensorflow/compiler:iree-import-xla) + list(APPEND _executable_paths integrations/tensorflow/compiler/iree-import-xla) +endif() + iree_add_bazel_invocation( INVOCATION_TARGET integrations_iree_tensorflow_importers - BAZEL_TARGETS - //integrations/tensorflow/compiler:iree-tf-import - EXECUTABLE_PATHS - integrations/tensorflow/compiler/iree-tf-import + BAZEL_TARGETS ${_bazel_targets} + EXECUTABLE_PATHS ${_executable_paths} ) if(${IREE_BUILD_PYTHON_BINDINGS})
diff --git a/integrations/tensorflow/bindings/python/CMakeLists.txt b/integrations/tensorflow/bindings/python/CMakeLists.txt index bafd98b..44fce9b 100644 --- a/integrations/tensorflow/bindings/python/CMakeLists.txt +++ b/integrations/tensorflow/bindings/python/CMakeLists.txt
@@ -19,6 +19,14 @@ add_subdirectory(${dir} "${_MAIN_PYTHON_DIR}/${dir}") endfunction() -_add_overlay_subdirectory(pyiree/tools/tf) +if(${IREE_BUILD_TENSORFLOW_COMPILER}) + _add_overlay_subdirectory(pyiree/tools/tf) +endif() + +if(${IREE_BUILD_XLA_COMPILER}) + _add_overlay_subdirectory(pyiree/tools/xla) +endif() + # TODO: Find another place for the TF support library. +# Pure python so can just always be generated. _add_overlay_subdirectory(pyiree/tf/support)
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/xla/CMakeLists.txt b/integrations/tensorflow/bindings/python/pyiree/tools/xla/CMakeLists.txt new file mode 100644 index 0000000..bcf4285 --- /dev/null +++ b/integrations/tensorflow/bindings/python/pyiree/tools/xla/CMakeLists.txt
@@ -0,0 +1,28 @@ +# Copyright 2020 Google LLC +# +# 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 +# +# https://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. + +iree_py_library( + NAME + xla + SRCS + "__init__.py" + DEPS + integrations_iree_tensorflow_importers +) + +iree_symlink_tool( + TARGET xla + FROM_TOOL_TARGET integrations_tensorflow_compiler_iree-import-xla + TO_EXE_NAME iree-import-xla +)
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/xla/__init__.py b/integrations/tensorflow/bindings/python/pyiree/tools/xla/__init__.py new file mode 100644 index 0000000..ceb37ea --- /dev/null +++ b/integrations/tensorflow/bindings/python/pyiree/tools/xla/__init__.py
@@ -0,0 +1,29 @@ +# Lint-as: python3 +"""TensorFlow tools.""" + +# Copyright 2020 Google LLC +# +# 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 +# +# https://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. + +from typing import Optional + +import os +import platform + + +def get_tool(exe_name: str) -> Optional[str]: + if platform.system() == "Windows": + exe_name = exe_name + ".exe" + this_path = os.path.dirname(__file__) + tool_path = os.path.join(this_path, exe_name) + return tool_path
diff --git a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/__init__.py b/integrations/tensorflow/bindings/python/pyiree/xla/compiler/__init__.py deleted file mode 100644 index fcfabe2..0000000 --- a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/__init__.py +++ /dev/null
@@ -1,108 +0,0 @@ -# Lint-as: python3 -"""Module init for the python bindings.""" - -# Copyright 2020 Google LLC -# -# 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 -# -# https://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. - -# pylint: disable=g-multiple-import -# pylint: disable=g-bad-import-order -# pylint: disable=g-import-not-at-top -# pylint: disable=wildcard-import - -__all__ = [ - # Common - "Context", - "Module", - "CompileOptions", - "OutputFormat", - # XLA - "XLA_IMPORT_PASS_PIPELINE", - "xla_load_module_proto", - "xla_compile_module_proto", -] - -from typing import Collection, Optional, Sequence - -from . import binding - -# Native aliases (matches those in the generic compiler). -llvm = binding.llvm -Context = binding.CompilerContext -Module = binding.CompilerModule -CompileOptions = binding.CompileOptions -OutputFormat = binding.OutputFormat - -# Pass pipeline that should run to lower a XLA-HLO module to a form suitable -# for input to the IREE compiler. -XLA_IMPORT_PASS_PIPELINE = ( - # Legalize to XLA - "canonicalize",) - - -# TODO(suderman): Update PyType to check the xla_computation is an XLA builder. -def xla_load_module_proto( - xla_computation, - compiler_context: Optional[Context] = None, - exported_names: Sequence[str] = (), - pass_pipeline: Sequence[str] = XLA_IMPORT_PASS_PIPELINE) -> Module: - """Loads a XLA saved model from its persistent representation. - - See also xla_compile_module_proto() for a one-shot API to load and compile. - - Args: - xla_computation: XLA Computation generate from XLA Python client - compiler_context: The pyiree.compiler.Context() backing the module. - exported_names: Optional sequence representing the exported names to keep. - pass_pipeline: Passes to run on the imported module prior to returning. - Defaults to XLA_IMPORT_PASS_PIPELINE. - - Returns: - An MLIR Module suitable for compilation by the IREE compiler. - This can be further compiled to an IREE blob by calling - .compile_to_sequencer_blob. - """ - if not compiler_context: - compiler_context = Context() - input_module = binding.load_xla_module_proto(compiler_context, - xla_computation, - exported_names=exported_names) - if pass_pipeline: - input_module.run_pass_pipeline(pass_pipeline) - return input_module - - -def xla_compile_module_proto( - xla_computation, - compiler_context: Optional[Context] = None, - exported_names: Sequence[str] = (), - pass_pipeline: Sequence[str] = XLA_IMPORT_PASS_PIPELINE, - target_backends: Sequence[str] = () -) -> binding.OpaqueBlob: - """Loads and compiles a XLA saved model in one shot. - - Args: - xla_computation: XLA Computation generate from XLA Python client - compiler_context: The pyiree.compiler.Context() backing the module. - exported_names: Optional sequence representing the exported names to keep. - pass_pipeline: Passes to run on the imported module prior to returning. - Defaults to XLA_IMPORT_PASS_PIPELINE. - target_backends: Optional sequence of specific target backends to compile - for (defaults to all compiled in targets). - - Returns: - An OpaqueBlob representing the compiled module. - """ - input_module = xla_load_module_proto(xla_computation, compiler_context, - exported_names, pass_pipeline) - return input_module.compile(target_backends=target_backends)
diff --git a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/export.def b/integrations/tensorflow/bindings/python/pyiree/xla/compiler/export.def deleted file mode 100644 index 1f2a8c1..0000000 --- a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/export.def +++ /dev/null
@@ -1,17 +0,0 @@ -;; Copyright 2020 Google LLC -;; -;; 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 -;; -;; https://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. - -LIBRARY BINDING -EXPORTS - PyInit_binding @1
diff --git a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/initialize_module.cc b/integrations/tensorflow/bindings/python/pyiree/xla/compiler/initialize_module.cc deleted file mode 100644 index 0b67c99..0000000 --- a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/initialize_module.cc +++ /dev/null
@@ -1,31 +0,0 @@ -// Copyright 2020 Google LLC -// -// 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 -// -// https://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 <mutex> // NOLINT - -#include "bindings/python/pyiree/common/binding.h" -#include "bindings/python/pyiree/compiler/compiler.h" -#include "integrations/tensorflow/bindings/python/pyiree/xla/compiler/register_xla.h" - -namespace iree { -namespace python { - -PYBIND11_MODULE(binding, m) { - m.doc() = "IREE XLA Compiler Interface"; - SetupCommonCompilerBindings(m); - SetupXlaBindings(m); -} - -} // namespace python -} // namespace iree
diff --git a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/register_xla.cc b/integrations/tensorflow/bindings/python/pyiree/xla/compiler/register_xla.cc deleted file mode 100644 index d2261aa..0000000 --- a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/register_xla.cc +++ /dev/null
@@ -1,69 +0,0 @@ -// Copyright 2020 Google LLC -// -// 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 -// -// https://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 "integrations/tensorflow/bindings/python/pyiree/xla/compiler/register_xla.h" - -#include <string> -#include <vector> - -#include "bindings/python/pyiree/common/status_utils.h" -#include "bindings/python/pyiree/compiler/compiler.h" -#include "llvm/Support/raw_ostream.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/MLIRContext.h" -#include "tensorflow/compiler/mlir/xla/hlo_function_importer.h" -#include "tensorflow/compiler/mlir/xla/hlo_module_importer.h" -#include "tensorflow/compiler/xla/client/xla_computation.h" - -namespace iree { -namespace python { - -namespace { - -CompilerModuleBundle LoadXlaModuleProto( - std::shared_ptr<CompilerContextBundle> context_bundle, - xla::XlaComputation& computation, - const std::vector<std::string>& exported_names) { - auto context = context_bundle->mlir_context(); - mlir::Builder builder(context); - - // TODO(suderman): Figure out how to transport a location detail in. - mlir::OwningModuleRef module = - mlir::ModuleOp::create(mlir::UnknownLoc::get(context)); - - xla::HloModuleImporter importer(module.get()); - auto result = importer.Import(computation.proto()); - if (!result.ok()) { - std::stringstream msg; - msg << "Failed to convert HLO Module Proto to MLIR: " << result; - throw RaisePyError(PyExc_RuntimeError, msg.str().c_str()); - } - - for (auto func : module->getOps<mlir::FuncOp>()) { - func.setAttr("iree.module.export", mlir::UnitAttr::get(func.getContext())); - } - - return CompilerModuleBundle(context_bundle, module.release()); -} - -} // namespace - -void SetupXlaBindings(pybind11::module m) { - m.def("load_xla_module_proto", &LoadXlaModuleProto, - py::arg("compiler_context"), py::arg("computation"), - py::arg("exported_names") = std::vector<std::string>()); -} - -} // namespace python -} // namespace iree
diff --git a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/register_xla.h b/integrations/tensorflow/bindings/python/pyiree/xla/compiler/register_xla.h deleted file mode 100644 index a2472ef..0000000 --- a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/register_xla.h +++ /dev/null
@@ -1,30 +0,0 @@ -// Copyright 2020 Google LLC -// -// 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 -// -// https://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. - -#ifndef IREE_BINDINGS_PYTHON_PYIREE_TF_INTEROP_REGISTER_XLA_H_ -#define IREE_BINDINGS_PYTHON_PYIREE_TF_INTEROP_REGISTER_XLA_H_ - -#include <string> - -#include "bindings/python/pyiree/common/binding.h" - -namespace iree { -namespace python { - -void SetupXlaBindings(pybind11::module m); - -} // namespace python -} // namespace iree - -#endif // IREE_BINDINGS_PYTHON_PYIREE_TF_INTEROP_REGISTER_XLA_H_
diff --git a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/xla_module_proto_test.py b/integrations/tensorflow/bindings/python/pyiree/xla/compiler/xla_module_proto_test.py deleted file mode 100644 index 02aeb39..0000000 --- a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/xla_module_proto_test.py +++ /dev/null
@@ -1,51 +0,0 @@ -# Copyright 2020 Google LLC -# -# 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 -# -# https://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. - -from absl.testing import absltest -import numpy as np -from pyiree.xla import compiler - -# pylint: disable=g-direct-tensorflow-import -import tensorflow.compiler.xla.python.xla_client as xla_client -# pylint: enable=g-direct-tensorflow-import - -ops = xla_client.ops - - -class RuntimeTest(absltest.TestCase): - - def testXLA(self): - """Tests that a basic saved model to XLA workflow grossly functions. - - This is largely here to verify that everything is linked in that needs to be - and that there are not no-ops, etc. - """ - # Generate a sample XLA computation. - builder = xla_client.XlaBuilder("testbuilder") - in_shape = np.array([4], dtype=np.float32) - in_feed = ops.Parameter(builder, 0, xla_client.shape_from_pyval(in_shape)) - result = ops.Add(in_feed, ops.Constant(builder, np.float32(1.0))) - xla_computation = builder.Build(result) - - # Load into XLA Module. - module = compiler.xla_load_module_proto(xla_computation) - - # Validate imported ASM. - xla_asm = module.to_asm() - print("XLA ASM: ", xla_asm) - self.assertRegex(xla_asm, "mhlo.add") - - -if __name__ == "__main__": - absltest.main()
diff --git a/integrations/tensorflow/compiler/BUILD b/integrations/tensorflow/compiler/BUILD index 6c91969..6ac0eca 100644 --- a/integrations/tensorflow/compiler/BUILD +++ b/integrations/tensorflow/compiler/BUILD
@@ -70,7 +70,7 @@ cc_binary( name = "iree-tf-opt", - srcs = ["iree-tf-opt-main.cc"], + srcs = ["iree-tf-opt-main.cpp"], deps = [ ":tensorflow", "//integrations/tensorflow/compiler/dialect/tf_strings/ir:dialect", @@ -98,7 +98,7 @@ cc_binary( name = "iree-tf-import", - srcs = ["iree-tf-import-main.cc"], + srcs = ["iree-tf-import-main.cpp"], deps = [ ":tensorflow", "@llvm-project//llvm:Support", @@ -118,3 +118,16 @@ ], }), ) + +cc_binary( + name = "iree-import-xla", + srcs = ["iree-import-xla-main.cpp"], + deps = [ + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", + "@org_tensorflow//tensorflow/compiler/mlir/xla:hlo_to_mlir_hlo", + "@org_tensorflow//tensorflow/compiler/xla/service:hlo_proto_cc", + "@org_tensorflow//tensorflow/core:lib", + ], +)
diff --git a/integrations/tensorflow/compiler/iree-import-xla-main.cpp b/integrations/tensorflow/compiler/iree-import-xla-main.cpp new file mode 100644 index 0000000..f64ac73 --- /dev/null +++ b/integrations/tensorflow/compiler/iree-import-xla-main.cpp
@@ -0,0 +1,178 @@ +// Copyright 2020 Google LLC +// +// 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 +// +// https://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. + +// Main entry-point for the XLA proto importer frontend. +// This will read from XLA proto files and produce MLIR MHLO assembly. + +#include <fstream> +#include <iostream> + +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/InitLLVM.h" +#include "llvm/Support/ToolOutputFile.h" +#include "mlir/IR/AsmState.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/Support/FileUtilities.h" +#include "tensorflow/compiler/mlir/xla/hlo_to_mlir_hlo.h" +#include "tensorflow/compiler/xla/service/hlo.pb.h" +#include "tensorflow/core/platform/protobuf.h" + +using namespace llvm; +using namespace mlir; + +namespace { + +enum XlaFormat { + binary_proto, + text_proto, +}; + +// Error collector that prints errors. +class PrintErrorCollector : public tensorflow::protobuf::io::ErrorCollector { + public: + PrintErrorCollector(std::string filePath) : filePath(std::move(filePath)) {} + void AddError(int line, int column, const std::string &message) override { + llvm::errs() << "Text protobuf parse error(" << filePath << ":" << line + << ":" << column << "): " << message << "\n"; + hadError = true; + } + + std::string filePath; + bool hadError = false; +}; + +class IStreamCopyingInputStream + : public tensorflow::protobuf::io::CopyingInputStream { + public: + IStreamCopyingInputStream(std::istream *input) : input(input) {} + int Read(void *buffer, int size) override { + input->read(static_cast<char *>(buffer), size); + if (input->fail()) return -1; + return input->gcount(); + } + + private: + std::istream *input; +}; + +} // namespace + +int main(int argc, char **argv) { + llvm::InitLLVM y(argc, argv); + + static cl::opt<std::string> inputPath( + cl::Positional, cl::desc("<XLA Protocol Buffer Path>"), cl::Required); + static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"), + cl::value_desc("filename"), + cl::init("-")); + static llvm::cl::opt<std::string> saveTempIreeImport( + "save-temp-iree-input", + llvm::cl::desc("Save the resultant IR to this file (useful for saving an " + "intermediate in a pipeline)"), + llvm::cl::init("")); + static llvm::cl::opt<XlaFormat> inputFormat( + "xla-format", cl::desc("XLA Format"), + cl::values(clEnumVal(binary_proto, "Parse a binary protocol buffer"), + clEnumVal(text_proto, "Parse a text protocol buffer"))); + + // Register any command line options. + registerAsmPrinterCLOptions(); + registerMLIRContextCLOptions(); + cl::ParseCommandLineOptions(argc, argv); + + DialectRegistry registry; + + // Read the protocol buffer. + std::ifstream fileInputStream; + std::istream *inputStream; + if (inputPath == "-") { + inputStream = &std::cin; + } else { + fileInputStream.open(inputPath, std::ios::in | std::ios::binary); + if (!fileInputStream.is_open()) { + llvm::errs() << "Unable to open input file " << inputPath << "\n"; + return 1; + } + inputStream = &fileInputStream; + } + + xla::HloProto hloProto; + switch (inputFormat) { + case binary_proto: { + if (!hloProto.mutable_hlo_module()->ParseFromIstream(inputStream)) { + llvm::errs() << "Could not parse binary protocol buffer from " + << inputPath << "\n"; + return 1; + } + break; + } + case text_proto: { + tensorflow::protobuf::TextFormat::Parser parser; + PrintErrorCollector collector(inputPath); + IStreamCopyingInputStream copyingStream(inputStream); + tensorflow::protobuf::io::CopyingInputStreamAdaptor streamAdaptor( + ©ingStream); + parser.RecordErrorsTo(&collector); + parser.Parse(&streamAdaptor, hloProto.mutable_hlo_module()); + if (collector.hadError) { + llvm::errs() << "Unable to parse text format protocol buffer\n"; + return 1; + } + break; + } + default: + llvm_unreachable("illegal XlaFormat"); + } + + // Convert the Module proto into MLIR. + MLIRContext context; + OwningModuleRef module = ModuleOp::create(mlir::UnknownLoc::get(&context)); + registry.loadAll(&context); + + auto status = + ConvertHloToMlirHlo(module.get(), hloProto.mutable_hlo_module()); + if (!status.ok()) { + llvm::errs() << "Error converting HLO Module Proto to MLIR: " + << status.ToString() << "\n"; + return 2; + } + + // Save. + auto saveToFile = [&](llvm::StringRef savePath) -> LogicalResult { + auto outputFile = openOutputFile(savePath); + if (!outputFile) { + llvm::errs() << "Could not open output file: " << savePath << "\n"; + return failure(); + } + OpPrintingFlags printFlags; + printFlags.enableDebugInfo(); + module->print(outputFile->os(), printFlags); + outputFile->os() << "\n"; + outputFile->keep(); + return success(); + }; + + // Save temp output. + if (!saveTempIreeImport.empty()) { + if (failed(saveToFile(saveTempIreeImport))) return 10; + } + + // Save output. + if (failed(saveToFile(outputFilename))) return 3; + return 0; +}
diff --git a/integrations/tensorflow/compiler/iree-tf-import-main.cc b/integrations/tensorflow/compiler/iree-tf-import-main.cpp similarity index 100% rename from integrations/tensorflow/compiler/iree-tf-import-main.cc rename to integrations/tensorflow/compiler/iree-tf-import-main.cpp
diff --git a/integrations/tensorflow/compiler/iree-tf-opt-main.cc b/integrations/tensorflow/compiler/iree-tf-opt-main.cpp similarity index 100% rename from integrations/tensorflow/compiler/iree-tf-opt-main.cc rename to integrations/tensorflow/compiler/iree-tf-opt-main.cpp
diff --git a/scripts/check_tabs.sh b/scripts/check_tabs.sh index f287fad..494ec10 100755 --- a/scripts/check_tabs.sh +++ b/scripts/check_tabs.sh
@@ -22,9 +22,10 @@ BASE_REF="${1:-main}" declare -a excluded_files_patterns=( - "^.gitmodules$" + "^\.gitmodules$" "/third_party/" "^third_party/" + "\.pb$" ) # Join on |