Merge pull request #4120 from google:main-to-google
PiperOrigin-RevId: 346216570
diff --git a/CMakeLists.txt b/CMakeLists.txt
index b37f0e0..9f08720 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -46,6 +46,7 @@
option(IREE_BUILD_PYTHON_BINDINGS "Builds the IREE python bindings" 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)
set(IREE_HAL_DRIVERS_TO_BUILD "all"
CACHE STRING "Semicolon-separated list of HAL drivers to build, or \"all\".")
@@ -268,6 +269,18 @@
include(iree_setup_toolchain)
#-------------------------------------------------------------------------------
+# Configure python early if there are any features that need it.
+# Note that doing this early ensures that dependencies that make incidental
+# use of Python (such as LLVM) resolve the same version.
+#-------------------------------------------------------------------------------
+
+if(${IREE_BUILD_COMPILER} OR
+ ${IREE_BUILD_PYTHON_BINDINGS} OR
+ ${IREE_BUILD_TENSORFLOW_COMPILER})
+ find_package(Python3 COMPONENTS Interpreter REQUIRED)
+endif()
+
+#-------------------------------------------------------------------------------
# MLIR/LLVM Dependency
# We treat the LLVM dependency specially because we support several different
# ways to use it:
@@ -371,22 +384,28 @@
endif()
#-------------------------------------------------------------------------------
-# Non-LLVM Dependencies
+# Python bindings.
#-------------------------------------------------------------------------------
-# Use the FindPython functions before any of our dependencies do. See
-# https://pybind11.readthedocs.io/en/stable/faq.html#inconsistent-detection-of-python-version-in-cmake-and-pybind11
-# If one dependency finds Python 2 (the default),
-# any others that try to find Python 3 will fail.
-# (Also come on, it's $CURRENT_YEAR - please just use Python 3 already.)
-if(${IREE_BUILD_COMPILER} OR ${IREE_BUILD_PYTHON_BINDINGS})
- find_package(Python3 COMPONENTS Interpreter REQUIRED)
-endif()
if(${IREE_BUILD_PYTHON_BINDINGS})
# Note: Optional because python libs can be manually specified.
find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
endif()
+#-------------------------------------------------------------------------------
+# Bazel setup (conditional on whether features need it)
+# Depends on python configuration.
+#-------------------------------------------------------------------------------
+
+if(${IREE_BUILD_TENSORFLOW_COMPILER})
+ include(configure_bazel)
+ iree_configure_bazel()
+endif()
+
+#-------------------------------------------------------------------------------
+# Other dependencies.
+#-------------------------------------------------------------------------------
+
include(external_cc_library)
include(flatbuffer_c_library)
@@ -516,6 +535,10 @@
add_subdirectory(experimental)
endif()
+if(${IREE_BUILD_TENSORFLOW_COMPILER})
+ add_subdirectory(integrations/tensorflow)
+endif()
+
if(${IREE_BUILD_PYTHON_BINDINGS})
iree_complete_py_extension_link_options()
endif()
diff --git a/bindings/python/CMakeLists.txt b/bindings/python/CMakeLists.txt
index 3febc3b..90a844c 100644
--- a/bindings/python/CMakeLists.txt
+++ b/bindings/python/CMakeLists.txt
@@ -17,4 +17,24 @@
set(PYBIND_COPTS "-fexceptions")
set(PYBIND_EXTENSION_COPTS "-fvisibility=hidden")
-add_subdirectory(pyiree)
+# Generated setup scripts.
+# TODO: Make the version configurable.
+set(IREE_PYTHON_VERSION "0.1a1")
+configure_file(setup.py setup.py COPYONLY)
+configure_file(setup_compiler.py.in setup_compiler.py)
+configure_file(setup_runtime.py.in setup_runtime.py)
+configure_file(setup_tools_core.py.in setup_tools_core.py)
+configure_file(setup_tools_tf.py.in setup_tools_tf.py)
+
+# Namespace packages.
+add_subdirectory(pyiree/common) # Deprecated
+add_subdirectory(pyiree/compiler2)
+add_subdirectory(pyiree/rt)
+
+if(${IREE_BUILD_COMPILER})
+add_subdirectory(pyiree/compiler) # Deprecated
+add_subdirectory(pyiree/tools/core)
+endif()
+
+# Tests.
+add_subdirectory(tests)
diff --git a/bindings/python/README.md b/bindings/python/README.md
index 4307803..18b7b44 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -1,19 +1,43 @@
-# IREE Python Sandbox
+# IREE Python API
-This directory contains various integration-oriented Python utilities that are
-not intended to be a public API. They are, however, useful for lower level
-compiler interop work. And of course, they are useful since we presently lack a
-real API :)
+Top-level packages:
-We're still untangling build support, jupyter integration, etc for OSS builds.
-Stand by.
+* `pyiree.compiler2` : Main compiler API (soon to be renamed to 'compiler').
+* `pyiree.rt` : Runtime components for executing binaries.
+* `pyiree.tools.core` : Core tools for executing the compiler.
+* `pyiree.tools.tf` : TensorFlow compiler tools (if enabled).
-## Issues:
+Deprecated packages:
-* This is called `pyiree` vs `iree` to avoid pythonpath collisions that tend
- to arise when an iree directory is inside of an iree directory.
-* The above could be solved in the bazel build by making iree/bindings/python
- its own sub-workspace.
-* However, doing so presently breaks both flatbuffer and tablegen generation
- because of fixes needed to those build rules so that they are sub-worksapce
- aware.
+* `pyiree.compiler`
+* `pyiree.common`
+* `pyiree.tf.compiler`
+
+## Installing
+
+First perform a normal CMake build with the following options:
+
+* `-DIREE_BUILD_PYTHON_BINDINGS=ON` : Enables Python Bindings
+* `-DIREE_BUILD_TENSORFLOW_COMPILER=ON` (optional) : Enables building the
+ TensorFlow compilers (note: requires additional dependencies. see overall
+ build docs).
+
+Then from the build directory, run:
+
+```shell
+# Install into a local installation or virtualenv.
+python bindings/python/setup.py install
+
+# Build wheels.
+python bindings/python/setup.py bdist_wheel
+```
+
+## Development mode
+
+For development, just set your `PYTHONPATH` environment variable to the
+`bindings/python` directory in your CMake build dir.
+
+## Run tests
+
+Tests under `bindings/python/tests` can be run directly once installed.
+Additional tests under `integrations/tensorflow/e2e` will be runnable soon.
diff --git a/bindings/python/pyiree/compiler2/.skip_bazel_to_cmake b/bindings/python/pyiree/compiler2/.skip_bazel_to_cmake
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/bindings/python/pyiree/compiler2/.skip_bazel_to_cmake
diff --git a/bindings/python/pyiree/CMakeLists.txt b/bindings/python/pyiree/compiler2/CMakeLists.txt
similarity index 76%
copy from bindings/python/pyiree/CMakeLists.txt
copy to bindings/python/pyiree/compiler2/CMakeLists.txt
index 55ac43d..c9a4ef8 100644
--- a/bindings/python/pyiree/CMakeLists.txt
+++ b/bindings/python/pyiree/compiler2/CMakeLists.txt
@@ -12,9 +12,15 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-add_subdirectory(common)
-add_subdirectory(rt)
+# Static and generated files.
+configure_file(README.md README.md COPYONLY)
-if(${IREE_BUILD_COMPILER})
- add_subdirectory(compiler)
-endif()
+iree_py_library(
+ NAME
+ compiler2
+ SRCS
+ "__init__.py"
+ "core.py"
+ "tf.py"
+ "tools.py"
+)
diff --git a/bindings/python/pyiree/compiler2/README.md b/bindings/python/pyiree/compiler2/README.md
new file mode 100644
index 0000000..f9d0e5b
--- /dev/null
+++ b/bindings/python/pyiree/compiler2/README.md
@@ -0,0 +1,46 @@
+# IREE Compiler Python Bindings
+
+Transitional note: These bindings are not complete yet and will ultimately
+replace the `pyiree.compiler` and `pyiree.tf.compiler` packages.
+
+## Core compiler
+
+```py
+from pyiree.compiler2 import *
+
+SIMPLE_MUL_ASM = """
+func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
+ attributes { iree.module.export } {
+ %0 = "mhlo.multiply"(%arg0, %arg1) {name = "mul.1"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
+ return %0 : tensor<4xf32>
+}
+"""
+
+# Also see compile_file()
+# There are many keyword options available.
+# See pyiree.compiler2.CompilerOptions
+binary = compile_str(SIMPLE_MUL_ASM, target_backends=["vulkan-spirv"])
+```
+
+
+## TensorFlow compiler
+
+```py
+import tensorflow as tf
+from pyiree.compiler2.tf import *
+
+class SimpleArithmeticModule(tf.Module):
+
+ @tf.function(input_signature=[
+ tf.TensorSpec([4], tf.float32),
+ tf.TensorSpec([4], tf.float32)
+ ])
+ def simple_mul(self, a, b):
+ return a * b
+
+# Also see compile_saved_model to directly compile an on-disk saved model.
+# There are many keyword options available.
+# See: pyiree.compiler2.tf.ImportOptions
+binary = compile_module(
+ SimpleArithmeticModule(), target_backends=["vulkan-spirv"])
+```
diff --git a/packaging/python/dummy_exclude_from_package.py b/bindings/python/pyiree/compiler2/__init__.py
similarity index 87%
rename from packaging/python/dummy_exclude_from_package.py
rename to bindings/python/pyiree/compiler2/__init__.py
index 0ca47f8..7ba6cb7 100644
--- a/packaging/python/dummy_exclude_from_package.py
+++ b/bindings/python/pyiree/compiler2/__init__.py
@@ -1,3 +1,4 @@
+# Lint-as: python3
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,3 +12,6 @@
# 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 .core import *
+from .tools import CompilerToolError
diff --git a/bindings/python/pyiree/compiler2/core.py b/bindings/python/pyiree/compiler2/core.py
new file mode 100644
index 0000000..cdc7900
--- /dev/null
+++ b/bindings/python/pyiree/compiler2/core.py
@@ -0,0 +1,191 @@
+# Lint-as: python3
+"""Core 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.
+
+from enum import Enum
+import subprocess
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+from .tools import *
+
+__all__ = [
+ "DEFAULT_TESTING_BACKENDS",
+ "compile_file",
+ "compile_str",
+ "CompilerOptions",
+ "OutputFormat",
+]
+
+# Default testing backend for invoking the compiler.
+DEFAULT_TESTING_BACKENDS = ["vmla"]
+
+
+class OutputFormat(Enum):
+ """The output format of the compiler."""
+ FLATBUFFER_BINARY = "flatbuffer-binary"
+ FLATBUFFER_TEXT = "flatbuffer-text"
+ MLIR_TEXT = "mlir-text"
+
+ @staticmethod
+ def parse(spec: Union[str, "OutputFormat"]) -> "OutputFormat":
+ """Parses or returns an OutputFormat.
+
+ Args:
+ spec: An OutputFormat instance or the case-insensitive name of one of
+ the enum values.
+ Returns:
+ An OutputFormat instance.
+ """
+ if isinstance(spec, OutputFormat):
+ return spec
+ spec = spec.upper()
+ if spec not in OutputFormat.__members__:
+ raise ValueError(f"For output_format= argument, expected one of: "
+ f"{', '.join(OutputFormat.__members__.keys())}")
+ return OutputFormat[spec]
+
+
+class CompilerOptions:
+ """Options to the compiler backend.
+
+ Keyword options:
+ output_file: Optionally save the compiled binary to a file instead of
+ returning it.
+ target_backends: List of str names of target backends to compile into
+ the binary. The resulting binary will run on targets that match one
+ or more of the compiled backends.
+ output_format: Override the output format. See the OutputFormat enum.
+ Values can either be an enum value or a case-insensitive name of
+ the option. Typically used for debugging
+ extra_args: Optional list of additional arguments to pass to the compiler.
+ Example: ["--print-ir-after-all"]
+ optimize: Whether to apply some default high level optimizations (default
+ True).
+ strip_debug_ops: Whether to strip high level operations used to aid
+ debugging.
+ strip_source_map: Whether to strip source map information (used to generate
+ better errors).
+ strip_symbols: Whether to strip extra symbols not needed for execution
+ (but which may aid debugging).
+ crash_reproducer_path: File name to output an MLIR crash dump to if there
+ is a compiler failure.
+ enable_benchmark: Whether to generate instrumented binaries suitable
+ for benchmarking.
+ """
+
+ def __init__(self,
+ *,
+ output_file: Optional[str] = None,
+ target_backends: Sequence[str] = (),
+ output_format: Union[OutputFormat,
+ str] = OutputFormat.FLATBUFFER_BINARY,
+ extra_args: Sequence[str] = (),
+ optimize: bool = True,
+ strip_debug_ops: bool = False,
+ strip_source_map: bool = False,
+ strip_symbols: bool = False,
+ crash_reproducer_path: Optional[str] = None,
+ enable_benchmark: bool = False):
+ self.output_file = output_file
+ self.target_backends = target_backends
+ self.output_format = OutputFormat.parse(output_format)
+ self.extra_args = extra_args
+ self.optimize = optimize
+ self.strip_debug_ops = strip_debug_ops
+ self.strip_source_map = strip_source_map
+ self.strip_symbols = strip_symbols
+ self.crash_reproducer_path = crash_reproducer_path
+ self.enable_benchmark = enable_benchmark
+
+
+def build_compile_command_line(input_file: str,
+ options: CompilerOptions) -> List[str]:
+ """Builds a command line for invoking the compiler.
+
+ Args:
+ input_file: The input file name.
+ options: Compiler options.
+ Returns:
+ List of strings of command line.
+ """
+ iree_translate = find_tool("iree-translate")
+ if not options.target_backends:
+ raise ValueError("Expected a non-empty list for 'target_backends'")
+
+ cl = [
+ iree_translate,
+ input_file,
+ f"--iree-vm-bytecode-module-output-format={options.output_format.value}",
+ f"--iree-hal-target-backends={','.join(options.target_backends)}",
+ ]
+
+ # Output file.
+ if options.output_file:
+ cl.append(f"-o={options.output_file}")
+
+ # Translation to perform.
+ cl.append("--iree-mlir-to-vm-bytecode-module" if not options.enable_benchmark
+ else "--iree-mlir-to-executable-benchmark-vm-module")
+
+ # Other options to set if specified.
+ if options.strip_debug_ops:
+ cl.append("--iree-vm-bytecode-module-strip-debug-ops")
+ if options.strip_source_map:
+ cl.append("--iree-vm-bytecode-module-strip-source-map")
+ if options.strip_symbols:
+ cl.append("--iree-vm-bytecode-module-strip-symbols")
+ if options.crash_reproducer_path:
+ cl.append(
+ f"--pass-pipeline-crash-reproducer={options.crash_reproducer_path}")
+
+ cl.extend(options.extra_args)
+ return cl
+
+
+def compile_file(input_file: str, **kwargs):
+ """Invokes the IREE compiler on an input file.
+
+ Args:
+ input_file: File containing MLIR assembly to compile.
+ **kwargs: Keyword arguments corresponding to CompilerOptions.
+ Returns:
+ Either a byte buffer of the compiled content or None if output_file
+ was specified in the options.
+ """
+ options = CompilerOptions(**kwargs)
+ cl = build_compile_command_line(input_file, options)
+ result = invoke_immediate(cl)
+ if options.output_file:
+ return None
+ return result
+
+
+def compile_str(input_str: str, **kwargs):
+ """Invokes the IREE compiler with an input string.
+
+ Args:
+ input_str: MLIR assembly to parse/compile.
+ **kwargs: Keyword arguments corresponding to CompilerOptions.
+ Returns:
+ Either a byte buffer of the compiled content or None if output_file
+ was specified in the options.
+ """
+ options = CompilerOptions(**kwargs)
+ cl = build_compile_command_line("-", options)
+ result = invoke_immediate(cl, immediate_input=input_str.encode("utf-8"))
+ if options.output_file:
+ return None
+ return result
diff --git a/bindings/python/pyiree/compiler2/setup.py b/bindings/python/pyiree/compiler2/setup.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/bindings/python/pyiree/compiler2/setup.py
diff --git a/bindings/python/pyiree/compiler2/tf.py b/bindings/python/pyiree/compiler2/tf.py
new file mode 100644
index 0000000..c9608af
--- /dev/null
+++ b/bindings/python/pyiree/compiler2/tf.py
@@ -0,0 +1,188 @@
+# Lint-as: python3
+"""TensorFlow 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.
+
+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_saved_model",
+ "compile_module",
+ "is_available",
+ "DEFAULT_TESTING_BACKENDS",
+ "ImportOptions",
+ "ImportType",
+]
+
+_TF_IMPORT_TOOL = "iree-tf-import"
+
+
+def is_available():
+ """Determine if TensorFlow and the compiler are available."""
+ try:
+ import tensorflow as tf
+ except ModuleNotFoundError:
+ logging.warn("Unable to import tensorflow")
+ return False
+ try:
+ find_tool(_TF_IMPORT_TOOL)
+ except ValueError:
+ logging.warning("Unable to find IREE tool %s", _TF_IMPORT_TOOL)
+ return False
+ return True
+
+
+class ImportType(Enum):
+ """Import type of the model."""
+ OBJECT_GRAPH = "savedmodel_v2"
+ V2 = "savedmodel_v2"
+ SIGNATURE_DEF = "savedmodel_v1"
+ V1 = "savedmodel_v1"
+
+ @staticmethod
+ def parse(spec: Union[str, "ImportType"]) -> "ImportType":
+ """Parses or returns an ImportType.
+
+ Args:
+ spec: An ImportType instance or the case-insensitive name of one of
+ the enum values.
+ Returns:
+ An ImportType instance.
+ """
+ if isinstance(spec, ImportType):
+ return spec
+ spec = spec.upper()
+ if spec not in ImportType.__members__:
+ raise ValueError(f"For import_type= argument, expected one of: "
+ f"{', '.join(ImportType.__members__.keys())}")
+ return ImportType[spec]
+
+
+class ImportOptions(CompilerOptions):
+ """Import options layer on top of the backend compiler options."""
+
+ def __init__(self,
+ exported_names: Sequence[str] = (),
+ import_only: bool = False,
+ import_type: Union[ImportType, str] = ImportType.OBJECT_GRAPH,
+ saved_model_tags: Set[str] = set(),
+ import_extra_args: Sequence[str] = (),
+ **kwargs):
+ """Initialize options from keywords.
+
+ Args:
+ exported_names: Optional sequence representing the exported names to
+ keep (object graph/v2 models only).
+ import_only: Only import the module. If True, the result will be textual
+ MLIR that can be further fed to the IREE compiler. If False (default),
+ the result will be the fully compiled IREE binary. In both cases,
+ bytes-like output is returned. Note that if the output_file= is
+ specified and import_only=True, then the MLIR form will be written to
+ the output file.
+ import_type: Type of import to perform. See ImportType enum.
+ saved_model_tags: Set of tags to export (signature def/v1 saved models
+ only).
+ import_extra_args: Extra arguments to pass to the iree-tf-import tool.
+ """
+ super().__init__(**kwargs)
+ self.exported_names = exported_names
+ self.import_only = import_only
+ self.import_type = ImportType.parse(import_type)
+ self.saved_model_tags = saved_model_tags
+ self.import_extra_args = import_extra_args
+
+
+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.
+ """
+ tf_import = find_tool(_TF_IMPORT_TOOL)
+ cl = [
+ tf_import,
+ input_path,
+ f"--tf-import-type={options.import_type.value}",
+ f"--tf-savedmodel-exported-names={','.join(options.exported_names)}",
+ f"--tf-savedmodel-tags={','.join(options.saved_model_tags)}",
+ ]
+ if options.import_only and options.output_file:
+ # Import stage directly outputs.
+ if options.output_file:
+ cl.append(f"-o={options.output_file}")
+ cl.extend(options.import_extra_args)
+ return cl
+
+
+def compile_saved_model(saved_model_dir: str, **kwargs):
+ """Compiles an on-disk saved model to an IREE binary.
+
+ Args:
+ saved_model_dir: Path to directory where the model was saved.
+ **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(saved_model_dir, 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_module(module, saved_model_dir: Optional[str] = None, **kwargs):
+ """Compiles a tf.Module to an IREE binary (by saving to disk).
+
+ Args:
+ module: The tf.Module instance to convert to MLIR
+ saved_model_dir: Optional path to save the tf.Module to. The module will not
+ be persisted on disk outside of this call if this is not provided.
+ **kwargs: Keyword args corresponding to ImportOptions or CompilerOptions.
+ Returns:
+ Same as compile_saved_model().
+ """
+
+ def do_it(saved_model_dir):
+ import tensorflow as tf
+ options = tf.saved_model.SaveOptions(save_debug_info=True)
+ tf.saved_model.save(module, saved_model_dir, options=options)
+ return compile_saved_model(saved_model_dir, **kwargs)
+
+ if saved_model_dir:
+ return do_it(saved_model_dir)
+ else:
+ with tempfile.TemporaryDirectory(suffix=".sm") as td:
+ return do_it(td)
diff --git a/bindings/python/pyiree/compiler2/tools.py b/bindings/python/pyiree/compiler2/tools.py
new file mode 100644
index 0000000..480600b
--- /dev/null
+++ b/bindings/python/pyiree/compiler2/tools.py
@@ -0,0 +1,254 @@
+# Lint-as: python3
+"""Utilities for locating and invoking compiler 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.
+
+import importlib
+import io
+import os
+import subprocess
+import sys
+import textwrap
+import threading
+
+from typing import List, Optional
+
+__all__ = [
+ "find_tool",
+ "invoke_immediate",
+ "invoke_pipeline",
+ "get_tool_path",
+ "CompilerToolError",
+]
+
+# In normal distribution circumstances, each named tool is associated with
+# 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-tf-import": "pyiree.tools.tf",
+ "iree-translate": "pyiree.tools.core",
+}
+
+# Map of tool module to package name as distributed to archives (used for
+# error messages).
+_TOOL_MODULE_PACKAGES = {
+ "pyiree.tools.core": "google-iree-tools-core",
+ "pyiree.tools.tf": "google-iree-tools-tf",
+}
+
+# Environment variable holding directories to be searched for named tools.
+# Delimitted by os.pathsep.
+_TOOL_PATH_ENVVAR = "IREE_TOOL_PATH"
+
+
+class CompilerToolError(Exception):
+ """Compiler exception that preserves the command line and error output."""
+
+ def __init__(self, process: subprocess.CompletedProcess):
+ try:
+ errs = process.stderr.decode("utf-8")
+ except:
+ errs = str(process.stderr) # Decode error or other: best we can do.
+
+ tool_name = os.path.basename(process.args[0])
+ super().__init__(f"Error invoking IREE compiler tool {tool_name}\n"
+ f"Diagnostics:\n{errs}\n\n"
+ f"Invoked with:\n {' '.join(process.args)}")
+
+
+def get_tool_path() -> List[str]:
+ """Returns list of paths to search for tools."""
+ list_str = os.environ.get(_TOOL_PATH_ENVVAR)
+ if not list_str:
+ return []
+ return list_str.split(os.pathsep)
+
+
+def find_tool(exe_name: str) -> str:
+ """Finds a tool by its (extension-less) executable name.
+
+ Args:
+ exe_name: The name of the executable (extension-less).
+ Returns:
+ An absolute path to the tool.
+ Raises:
+ ValueError: If the tool is not known or not found.
+ """
+ if exe_name not in _TOOL_MODULE_MAP:
+ raise ValueError(f"IREE compiler tool '{exe_name}' is not a known tool")
+ # First search an explicit tool path.
+ tool_path = get_tool_path()
+ for path_entry in tool_path:
+ if not path_entry:
+ continue
+ candidate_exe = os.path.join(path_entry, exe_name)
+ if os.path.isfile(candidate_exe) and os.access(candidate_exe, os.X_OK):
+ return candidate_exe
+
+ # Attempt to load the tool module.
+ tool_module_name = _TOOL_MODULE_MAP[exe_name]
+ tool_module_package = _TOOL_MODULE_PACKAGES[tool_module_name]
+ try:
+ tool_module = importlib.import_module(tool_module_name)
+ except ModuleNotFoundError:
+ raise ValueError(
+ f"IREE compiler tool '{exe_name}' is not installed (it should have been "
+ f"found in the python module '{tool_module_name}', typically installed "
+ f"via the package {tool_module_package}).\n\n"
+ f"Either install the package or set the {_TOOL_PATH_ENVVAR} environment "
+ f"variable to contain the path of the tool executable "
+ f"(current {_TOOL_PATH_ENVVAR} = {repr(tool_path)})") from None
+
+ # Ask the module for its tool.
+ candidate_exe = tool_module.get_tool(exe_name)
+ if (not candidate_exe or not os.path.isfile(candidate_exe) or
+ not os.access(candidate_exe, os.X_OK)):
+ raise ValueError(
+ f"IREE compiler tool '{exe_name}' was located in module "
+ f"'{tool_module_name}' but the file was not found or not executable: "
+ f"{candidate_exe}")
+ return candidate_exe
+
+
+def invoke_immediate(command_line: List[str],
+ *,
+ input_file: Optional[str] = None,
+ immediate_input=None):
+ """Invokes an immediate command.
+
+ This is separate from invoke_pipeline as it is simpler and supports more
+ complex input redirection, using recommended facilities for sub-processes
+ (less magic).
+
+ Note that this differs from the usual way of using subprocess.run or
+ subprocess.Popen().communicate() because we need to pump all of the error
+ streams individually and only pump pipes not connected to a different stage.
+ Uses threads to pump everything that is required.
+ """
+ run_args = {}
+ input_file_handle = None
+ stderr_handle = sys.stderr
+ try:
+ # Redirect input.
+ if input_file is not None:
+ input_file_handle = open(input_file, "rb")
+ run_args["stdin"] = input_file_handle
+ elif immediate_input is not None:
+ run_args["input"] = immediate_input
+
+ # Capture output.
+ # Upgrade note: Python >= 3.7 can just use capture_output=True
+ run_args["stdout"] = subprocess.PIPE
+ run_args["stderr"] = subprocess.PIPE
+ process = subprocess.run(command_line, **run_args)
+ if process.returncode != 0:
+ raise CompilerToolError(process)
+ # Emit stderr contents.
+ _write_binary_stderr(stderr_handle, process.stderr)
+ return process.stdout
+ finally:
+ if input_file_handle:
+ input_file_handle.close()
+
+
+def invoke_pipeline(command_lines: List[List[str]]):
+ """Invoke a pipeline of commands.
+
+ The first stage of the pipeline will have its stdin set to DEVNULL and each
+ subsequent stdin will derive from the prior stdout. The final stdout will
+ be accumulated and returned. All stderr contents are accumulated and printed
+ to stderr on completion or the first failing stage of the pipeline will have
+ an exception raised with its stderr output.
+ """
+ stages = []
+ prev_out = subprocess.DEVNULL
+ stderr_handle = sys.stderr
+
+ # Create all stages.
+ for i in range(len(command_lines)):
+ command_line = command_lines[i]
+ popen_args = {
+ "stdin": prev_out,
+ "stdout": subprocess.PIPE,
+ "stderr": subprocess.PIPE,
+ }
+ process = subprocess.Popen(command_line, **popen_args)
+ prev_out = process.stdout
+ capture_output = (i == (len(command_lines) - 1))
+ stages.append(_PipelineStage(process, capture_output))
+
+ # Start stages.
+ for stage in stages:
+ stage.start()
+
+ # Join.
+ for stage in stages:
+ stage.join()
+
+ # Check for errors.
+ for stage in stages:
+ assert stage.completed
+ if stage.completed.returncode != 0:
+ raise CompilerToolError(stage.completed)
+
+ # Print any stderr output.
+ for stage in stages:
+ _write_binary_stderr(stderr_handle, stage.errs)
+ return stages[-1].outs
+
+
+class _PipelineStage(threading.Thread):
+ """Wraps a process and pumps its handles, waiting for completion."""
+
+ def __init__(self, process, capture_output):
+ super().__init__()
+ self.process = process
+ self.capture_output = capture_output
+ self.completed: Optional[subprocess.CompletedProcess] = None
+ self.outs = None
+ self.errs = None
+
+ def pump_stderr(self):
+ self.errs = self.process.stderr.read()
+
+ def pump_stdout(self):
+ self.outs = self.process.stdout.read()
+
+ def run(self):
+ stderr_thread = threading.Thread(target=self.pump_stderr)
+ stderr_thread.start()
+ if self.capture_output:
+ stdout_thread = threading.Thread(target=self.pump_stdout)
+ stdout_thread.start()
+ self.process.wait()
+ stderr_thread.join()
+ if self.capture_output:
+ stdout_thread.join()
+ self.completed = subprocess.CompletedProcess(self.process.args,
+ self.process.returncode,
+ self.outs, self.errs)
+ self.process.stderr.close()
+ self.process.stdout.close()
+
+
+def _write_binary_stderr(out_handle, contents):
+ # Fast-paths buffered text-io (which stderr is by default) while allowing
+ # full decode for non buffered and binary io.
+ if hasattr(out_handle, "buffer"):
+ out_handle.buffer.write(contents)
+ elif isinstance(out_handle, io.TextIOBase):
+ out_handle.write(contents.decode("utf-8"))
+ else:
+ out_handle.write(contents)
diff --git a/bindings/python/pyiree/rt/CMakeLists.txt b/bindings/python/pyiree/rt/CMakeLists.txt
index 7828073..ef0f5bb 100644
--- a/bindings/python/pyiree/rt/CMakeLists.txt
+++ b/bindings/python/pyiree/rt/CMakeLists.txt
@@ -12,6 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+# Static and generated files.
+configure_file(README.md README.md COPYONLY)
+
iree_pyext_library(
NAME
PyExtRtLib
diff --git a/bindings/python/pyiree/rt/README.md b/bindings/python/pyiree/rt/README.md
new file mode 100644
index 0000000..44b94e1
--- /dev/null
+++ b/bindings/python/pyiree/rt/README.md
@@ -0,0 +1,4 @@
+# IREE Python Runtime Components
+
+This package provides an API for running compiled IREE binaries and interfacing
+with the hardware-abstraction-layer.
diff --git a/bindings/python/pyiree/tools/core/.skip_bazel_to_cmake b/bindings/python/pyiree/tools/core/.skip_bazel_to_cmake
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/bindings/python/pyiree/tools/core/.skip_bazel_to_cmake
diff --git a/bindings/python/pyiree/CMakeLists.txt b/bindings/python/pyiree/tools/core/CMakeLists.txt
similarity index 77%
copy from bindings/python/pyiree/CMakeLists.txt
copy to bindings/python/pyiree/tools/core/CMakeLists.txt
index 55ac43d..9e4364c 100644
--- a/bindings/python/pyiree/CMakeLists.txt
+++ b/bindings/python/pyiree/tools/core/CMakeLists.txt
@@ -12,9 +12,17 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-add_subdirectory(common)
-add_subdirectory(rt)
+iree_py_library(
+ NAME
+ core
+ SRCS
+ "__init__.py"
+)
if(${IREE_BUILD_COMPILER})
- add_subdirectory(compiler)
+ iree_symlink_tool(
+ TARGET core
+ FROM_TOOL_TARGET iree_tools_iree-translate
+ TO_EXE_NAME iree-translate
+ )
endif()
diff --git a/bindings/python/pyiree/CMakeLists.txt b/bindings/python/pyiree/tools/core/__init__.py
similarity index 64%
copy from bindings/python/pyiree/CMakeLists.txt
copy to bindings/python/pyiree/tools/core/__init__.py
index 55ac43d..4ce9c14 100644
--- a/bindings/python/pyiree/CMakeLists.txt
+++ b/bindings/python/pyiree/tools/core/__init__.py
@@ -1,3 +1,6 @@
+# Lint-as: python3
+"""Core tools."""
+
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,9 +15,15 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-add_subdirectory(common)
-add_subdirectory(rt)
+from typing import Optional
-if(${IREE_BUILD_COMPILER})
- add_subdirectory(compiler)
-endif()
+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/bindings/python/setup.py b/bindings/python/setup.py
new file mode 100644
index 0000000..3a669cc
--- /dev/null
+++ b/bindings/python/setup.py
@@ -0,0 +1,41 @@
+#!/usr/bin/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.
+
+# Build platform specific wheel files for the pyiree.rt package.
+# Built artifacts are per-platform and build out of the build tree.
+
+import os
+import subprocess
+import sys
+
+# Make this setup position independent and make it not conflict with
+# parallel scripts.
+this_dir = os.path.abspath(os.path.dirname(__file__))
+
+
+def run_sub_setup(name):
+ sub_path = os.path.join(this_dir, f"{name}.py")
+ args = [sys.executable, sub_path] + sys.argv[1:]
+ print(f"##### Running sub setup: {' '.join(args)}")
+ subprocess.check_call(args)
+ print("")
+
+
+run_sub_setup("setup_compiler")
+run_sub_setup("setup_runtime")
+run_sub_setup("setup_tools_core")
+if os.path.exists(os.path.join(this_dir, "pyiree/tools/tf")):
+ run_sub_setup("setup_tools_tf")
diff --git a/bindings/python/setup_compiler.py.in b/bindings/python/setup_compiler.py.in
new file mode 100644
index 0000000..f8d637d
--- /dev/null
+++ b/bindings/python/setup_compiler.py.in
@@ -0,0 +1,57 @@
+#!/usr/bin/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.
+
+# Build platform specific wheel files for the pyiree.rt package.
+# Built artifacts are per-platform and build out of the build tree.
+
+import os
+from setuptools import setup, find_namespace_packages
+
+# Make this setup position independent and make it not conflict with
+# parallel scripts.
+this_dir = os.path.abspath(os.path.dirname(__file__))
+setup_dir = os.path.join(this_dir, "setupbuild", "compiler")
+os.makedirs(setup_dir, exist_ok=True)
+os.chdir(setup_dir)
+
+def read(fname):
+ return open(os.path.join(this_dir, fname), "rt").read()
+
+
+setup(
+ name="google-iree-compiler",
+ version="@IREE_PYTHON_VERSION@",
+ author="The IREE Team",
+ author_email="iree-discuss@googlegroups.com",
+ license="Apache",
+ description="IREE Python Compiler API",
+ long_description=read("pyiree/compiler2/README.md"),
+ long_description_content_type="text/markdown",
+ url="https://github.com/google/iree",
+ classifiers=[
+ "Programming Language :: Python :: 3",
+ "License :: OSI Approved :: Apache License",
+ "Operating System :: OS Independent",
+ "Development Status :: 3 - Alpha",
+ ],
+ python_requires=">=3.6",
+ package_dir={"": this_dir},
+ packages=find_namespace_packages(
+ where=this_dir,
+ include=["pyiree.compiler2", "pyiree.compiler2.*"],
+ exclude=["*.CMakeFiles"]),
+ zip_safe=False, # This package is fine but not zipping is more versatile.
+)
diff --git a/bindings/python/setup_runtime.py.in b/bindings/python/setup_runtime.py.in
new file mode 100644
index 0000000..5893017
--- /dev/null
+++ b/bindings/python/setup_runtime.py.in
@@ -0,0 +1,66 @@
+#!/usr/bin/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.
+
+# Build platform specific wheel files for the pyiree.rt package.
+# Built artifacts are per-platform and build out of the build tree.
+
+import os
+from setuptools import setup, find_namespace_packages, Extension
+import sysconfig
+
+# Make this setup position independent and make it not conflict with
+# parallel scripts.
+this_dir = os.path.abspath(os.path.dirname(__file__))
+setup_dir = os.path.join(this_dir, "setupbuild", "runtime")
+os.makedirs(setup_dir, exist_ok=True)
+os.chdir(setup_dir)
+
+
+def read(fname):
+ return open(os.path.join(this_dir, fname), "rt").read()
+
+
+setup(
+ name="google-iree-runtime",
+ version="@IREE_PYTHON_VERSION@",
+ author="The IREE Team",
+ author_email="iree-discuss@googlegroups.com",
+ license="Apache",
+ description="IREE Python Runtime Components",
+ long_description=read("pyiree/rt/README.md"),
+ long_description_content_type="text/markdown",
+ url="https://github.com/google/iree",
+ classifiers=[
+ "Programming Language :: Python :: 3",
+ "License :: OSI Approved :: Apache License",
+ "Operating System :: OS Independent",
+ "Development Status :: 3 - Alpha",
+ ],
+ python_requires=">=3.6",
+ package_dir={"": this_dir},
+ packages=find_namespace_packages(where=this_dir,
+ include=["pyiree.rt", "pyiree.rt.*"],
+ exclude=["*.CMakeFiles"]),
+ ext_modules=[
+ Extension(name="pyiree.rt.binding", sources=[]),
+ ],
+ # Matching the native extension as a data file keeps setuptools from
+ # "building" it (i.e. turning it into a static binary).
+ package_data={
+ "": [f"*{sysconfig.get_config_var('EXT_SUFFIX')}"],
+ },
+ zip_safe=False,
+)
diff --git a/bindings/python/setup_tools_core.py.in b/bindings/python/setup_tools_core.py.in
new file mode 100644
index 0000000..55104f5
--- /dev/null
+++ b/bindings/python/setup_tools_core.py.in
@@ -0,0 +1,80 @@
+#!/usr/bin/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.
+
+# Build platform specific wheel files for the pyiree.rt package.
+# Built artifacts are per-platform and build out of the build tree.
+
+import os
+import platform
+from setuptools import setup, find_namespace_packages
+
+# Make this setup position independent and make it not conflict with
+# parallel scripts.
+this_dir = os.path.abspath(os.path.dirname(__file__))
+setup_dir = os.path.join(this_dir, "setupbuild", "tools_core")
+os.makedirs(setup_dir, exist_ok=True)
+os.chdir(setup_dir)
+
+exe_suffix = ".exe" if platform.system() == "Windows" else ""
+
+
+def read(fname):
+ return open(os.path.join(this_dir, fname), "rt").read()
+
+
+# Force platform specific wheel.
+# https://stackoverflow.com/questions/45150304
+try:
+ from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
+
+ class bdist_wheel(_bdist_wheel):
+
+ def finalize_options(self):
+ _bdist_wheel.finalize_options(self)
+ self.root_is_pure = False
+except ImportError:
+ bdist_wheel = None
+
+setup(
+ name="google-iree-tools-core",
+ version="@IREE_PYTHON_VERSION@",
+ author="The IREE Team",
+ author_email="iree-discuss@googlegroups.com",
+ license="Apache",
+ description="IREE Python Core Tools Binaries",
+ long_description=
+ "Package containing platform-specific binaries for core compiler tools",
+ long_description_content_type="text/plain",
+ url="https://github.com/google/iree",
+ classifiers=[
+ "Programming Language :: Python :: 3",
+ "License :: OSI Approved :: Apache License",
+ "Operating System :: OS Independent",
+ "Development Status :: 3 - Alpha",
+ ],
+ python_requires=">=3.6",
+ package_dir={"": this_dir},
+ packages=find_namespace_packages(where=this_dir,
+ include=["pyiree.tools.core"],
+ exclude=["*.CMakeFiles"]),
+ # Matching the native extension as a data file keeps setuptools from
+ # "building" it (i.e. turning it into a static binary).
+ package_data={
+ "pyiree.tools.core": [f"iree-translate{exe_suffix}",],
+ },
+ cmdclass={'bdist_wheel': bdist_wheel},
+ zip_safe=False,
+)
diff --git a/bindings/python/setup_tools_tf.py.in b/bindings/python/setup_tools_tf.py.in
new file mode 100644
index 0000000..a1b292e
--- /dev/null
+++ b/bindings/python/setup_tools_tf.py.in
@@ -0,0 +1,81 @@
+#!/usr/bin/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.
+
+# Build platform specific wheel files for the pyiree.rt package.
+# Built artifacts are per-platform and build out of the build tree.
+
+import os
+import platform
+from setuptools import setup, find_namespace_packages
+
+# Make this setup position independent and make it not conflict with
+# parallel scripts.
+this_dir = os.path.abspath(os.path.dirname(__file__))
+setup_dir = os.path.join(this_dir, "setupbuild", "tools_tf")
+os.makedirs(setup_dir, exist_ok=True)
+os.chdir(setup_dir)
+
+exe_suffix = ".exe" if platform.system() == "Windows" else ""
+
+
+def read(fname):
+ return open(os.path.join(this_dir, fname), "rt").read()
+
+
+# Force platform specific wheel.
+# https://stackoverflow.com/questions/45150304
+try:
+ from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
+
+ class bdist_wheel(_bdist_wheel):
+
+ def finalize_options(self):
+ _bdist_wheel.finalize_options(self)
+ self.root_is_pure = False
+except ImportError:
+ bdist_wheel = None
+
+setup(
+ name="google-iree-tools-tf",
+ version="@IREE_PYTHON_VERSION@",
+ author="The IREE Team",
+ author_email="iree-discuss@googlegroups.com",
+ license="Apache",
+ description="IREE Python TensorFlow Tools Binaries",
+ long_description=
+ "Package containing platform-specific binaries for TensorFlow "
+ "compiler tools",
+ long_description_content_type="text/plain",
+ url="https://github.com/google/iree",
+ classifiers=[
+ "Programming Language :: Python :: 3",
+ "License :: OSI Approved :: Apache License",
+ "Operating System :: OS Independent",
+ "Development Status :: 3 - Alpha",
+ ],
+ python_requires=">=3.6",
+ package_dir={"": this_dir},
+ packages=find_namespace_packages(where=this_dir,
+ include=["pyiree.tools.tf"],
+ exclude=["*.CMakeFiles"]),
+ # Matching the native extension as a data file keeps setuptools from
+ # "building" it (i.e. turning it into a static binary).
+ package_data={
+ "pyiree.tools.tf": [f"iree-tf-import{exe_suffix}",],
+ },
+ cmdclass={'bdist_wheel': bdist_wheel},
+ zip_safe=False,
+)
diff --git a/bindings/python/tests/.skip_bazel_to_cmake b/bindings/python/tests/.skip_bazel_to_cmake
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/bindings/python/tests/.skip_bazel_to_cmake
diff --git a/bindings/python/pyiree/CMakeLists.txt b/bindings/python/tests/CMakeLists.txt
similarity index 78%
rename from bindings/python/pyiree/CMakeLists.txt
rename to bindings/python/tests/CMakeLists.txt
index 55ac43d..701bb72 100644
--- a/bindings/python/pyiree/CMakeLists.txt
+++ b/bindings/python/tests/CMakeLists.txt
@@ -12,9 +12,16 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-add_subdirectory(common)
-add_subdirectory(rt)
+iree_py_test(
+ NAME
+ compiler_core_test
+ SRCS
+ "compiler_core_test.py"
+)
-if(${IREE_BUILD_COMPILER})
- add_subdirectory(compiler)
-endif()
+iree_py_test(
+ NAME
+ compiler_tf_test
+ SRCS
+ "compiler_tf_test.py"
+)
diff --git a/bindings/python/tests/README.md b/bindings/python/tests/README.md
new file mode 100644
index 0000000..a0a288e
--- /dev/null
+++ b/bindings/python/tests/README.md
@@ -0,0 +1,8 @@
+# Python API Tests
+
+These tests are run in an environment where all available Python bindings
+are setup on the `PYTHONPATH`. Each will internally skip itself if optional
+components are not available.
+
+Note that IREE compiler tool locations can be overriden by specifying the
+`IREE_TOOL_PATH` environment variable.
diff --git a/bindings/python/tests/compiler_core_test.py b/bindings/python/tests/compiler_core_test.py
new file mode 100644
index 0000000..54c15bd
--- /dev/null
+++ b/bindings/python/tests/compiler_core_test.py
@@ -0,0 +1,138 @@
+# 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 contextlib
+import logging
+import os
+import io
+import tempfile
+import unittest
+
+from pyiree import compiler2 as compiler
+
+SIMPLE_MUL_ASM = """
+func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
+ attributes { iree.module.export } {
+ %0 = "mhlo.multiply"(%arg0, %arg1) {name = "mul.1"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
+ return %0 : tensor<4xf32>
+}
+"""
+
+
+class CompilerTest(unittest.TestCase):
+
+ def testNoTargetBackends(self):
+ with self.assertRaisesRegex(
+ ValueError, "Expected a non-empty list for 'target_backends'"):
+ binary = compiler.compile_str(SIMPLE_MUL_ASM)
+
+ def testCompileStr(self):
+ binary = compiler.compile_str(
+ SIMPLE_MUL_ASM, target_backends=compiler.DEFAULT_TESTING_BACKENDS)
+ logging.info("Flatbuffer size = %d", len(binary))
+ self.assertTrue(binary)
+
+ def testCompileInputFile(self):
+ with tempfile.NamedTemporaryFile("wt", delete=False) as f:
+ try:
+ f.write(SIMPLE_MUL_ASM)
+ f.close()
+ binary = compiler.compile_file(
+ f.name, target_backends=compiler.DEFAULT_TESTING_BACKENDS)
+ finally:
+ os.remove(f.name)
+ logging.info("Flatbuffer size = %d", len(binary))
+ self.assertIn(b"simple_mul", binary)
+
+ def testCompileOutputFile(self):
+ with tempfile.NamedTemporaryFile("wt", delete=False) as f:
+ try:
+ f.close()
+ output = compiler.compile_str(
+ SIMPLE_MUL_ASM,
+ output_file=f.name,
+ target_backends=compiler.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("Flatbuffer size = %d", len(binary))
+ self.assertIn(b"simple_mul", binary)
+
+ def testOutputFbText(self):
+ text = compiler.compile_str(
+ SIMPLE_MUL_ASM,
+ output_format=compiler.OutputFormat.FLATBUFFER_TEXT,
+ target_backends=compiler.DEFAULT_TESTING_BACKENDS).decode("utf-8")
+ # Just check for an arbitrary JSON-tag.
+ self.assertIn('"exported_functions"', text)
+
+ def testBadOutputFormat(self):
+ with self.assertRaisesRegex(
+ ValueError, "For output_format= argument, expected one of: "
+ "FLATBUFFER_BINARY, FLATBUFFER_TEXT, MLIR_TEXT"):
+ _ = compiler.compile_str(
+ SIMPLE_MUL_ASM,
+ output_format="foobar",
+ target_backends=compiler.DEFAULT_TESTING_BACKENDS)
+
+ def testOutputFbTextParsed(self):
+ text = compiler.compile_str(
+ SIMPLE_MUL_ASM,
+ output_format='flatbuffer_text',
+ target_backends=compiler.DEFAULT_TESTING_BACKENDS).decode("utf-8")
+ # Just check for an arbitrary JSON-tag.
+ self.assertIn('"exported_functions"', text)
+
+ def testOutputMlirText(self):
+ text = compiler.compile_str(
+ SIMPLE_MUL_ASM,
+ output_format=compiler.OutputFormat.MLIR_TEXT,
+ target_backends=compiler.DEFAULT_TESTING_BACKENDS).decode("utf-8")
+ # Just check for a textual op name.
+ self.assertIn("vm.module", text)
+
+ def testExtraArgsStderr(self):
+ # pass-timing is not special: it just does something and emits to stderr.
+ with io.StringIO() as buf, contextlib.redirect_stderr(buf):
+ compiler.compile_str(SIMPLE_MUL_ASM,
+ extra_args=["--pass-timing"],
+ target_backends=compiler.DEFAULT_TESTING_BACKENDS)
+ stderr = buf.getvalue()
+ self.assertIn("Pass execution timing report", stderr)
+
+ def testAllOptions(self):
+ binary = compiler.compile_str(
+ SIMPLE_MUL_ASM,
+ optimize=False,
+ strip_debug_ops=True,
+ strip_source_map=True,
+ strip_symbols=True,
+ crash_reproducer_path="foobar.txt",
+ enable_benchmark=True,
+ target_backends=compiler.DEFAULT_TESTING_BACKENDS)
+
+ def testException(self):
+ with self.assertRaisesRegex(compiler.CompilerToolError, "Invoked with"):
+ _ = compiler.compile_str(
+ "I'm a little teapot but not a valid program",
+ target_backends=compiler.DEFAULT_TESTING_BACKENDS)
+
+
+if __name__ == "__main__":
+ logging.basicConfig(level=logging.DEBUG)
+ unittest.main()
diff --git a/bindings/python/tests/compiler_tf_test.py b/bindings/python/tests/compiler_tf_test.py
new file mode 100644
index 0000000..41fc94e
--- /dev/null
+++ b/bindings/python/tests/compiler_tf_test.py
@@ -0,0 +1,89 @@
+# 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.tf import *
+
+if not is_available():
+ print(f"Skipping test {__file__} because the IREE TensorFlow compiler "
+ f"is not installed")
+ sys.exit(0)
+
+import tensorflow as tf
+
+
+class SimpleArithmeticModule(tf.Module):
+
+ @tf.function(input_signature=[
+ tf.TensorSpec([4], tf.float32),
+ tf.TensorSpec([4], tf.float32)
+ ])
+ def simple_mul(self, a, b):
+ return a * b
+
+ @tf.function(input_signature=[
+ tf.TensorSpec([128, 3072], tf.float32),
+ tf.TensorSpec([3072, 256], tf.float32),
+ ])
+ def simple_matmul(self, a, b):
+ return tf.matmul(a, b)
+
+
+# TODO(laurenzo): More test cases needed (may need additional files).
+# Specifically, figure out how to test v1 models.
+class TfCompilerTest(unittest.TestCase):
+
+ def testImportSavedModel(self):
+ import_mlir = compile_saved_model(self.smdir,
+ import_only=True).decode("utf-8")
+ self.assertIn("func @simple_matmul", import_mlir)
+
+ def testCompileSavedModel(self):
+ binary = compile_saved_model(self.smdir,
+ target_backends=DEFAULT_TESTING_BACKENDS)
+ logging.info("Compiled len: %d", len(binary))
+ self.assertIn(b"simple_matmul", binary)
+ self.assertIn(b"simple_mul", binary)
+
+ def testCompileModule(self):
+ binary = compile_module(self.m, target_backends=DEFAULT_TESTING_BACKENDS)
+ logging.info("Compiled len: %d", len(binary))
+ self.assertIn(b"simple_matmul", binary)
+ self.assertIn(b"simple_mul", binary)
+
+ @classmethod
+ def setUpClass(cls):
+ cls.m = SimpleArithmeticModule()
+ cls.tempdir = tempfile.TemporaryDirectory()
+ cls.smdir = os.path.join(cls.tempdir.name, "arith.sm")
+ tf.saved_model.save(
+ cls.m,
+ cls.smdir,
+ options=tf.saved_model.SaveOptions(save_debug_info=True))
+
+ @classmethod
+ def tearDownClass(cls):
+ cls.tempdir.cleanup()
+
+
+if __name__ == "__main__":
+ logging.basicConfig(level=logging.DEBUG)
+ unittest.main()
diff --git a/build_tools/bazel/build_tensorflow.sh b/build_tools/bazel/build_tensorflow.sh
index cc9a407..055e586 100755
--- a/build_tools/bazel/build_tensorflow.sh
+++ b/build_tools/bazel/build_tensorflow.sh
@@ -84,7 +84,7 @@
--nosystem_rc --nohome_rc --noworkspace_rc \
--bazelrc=build_tools/bazel/iree.bazelrc \
query \
- //integrations/... + //colab/... + //packaging/... | \
+ //integrations/... + //colab/... | \
xargs bazel \
--nosystem_rc --nohome_rc --noworkspace_rc \
--bazelrc=build_tools/bazel/iree.bazelrc \
diff --git a/build_tools/cmake/bazel.bat.in b/build_tools/cmake/bazel.bat.in
new file mode 100644
index 0000000..93da790
--- /dev/null
+++ b/build_tools/cmake/bazel.bat.in
@@ -0,0 +1,17 @@
+@echo off
+REM Copyright 2020 Google LLC
+REM
+REM Licensed under the Apache License, Version 2.0 (the "License");
+REM you may not use this file except in compliance with the License.
+REM You may obtain a copy of the License at
+REM
+REM https://www.apache.org/licenses/LICENSE-2.0
+REM
+REM Unless required by applicable law or agreed to in writing, software
+REM distributed under the License is distributed on an "AS IS" BASIS,
+REM WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+REM See the License for the specific language governing permissions and
+REM limitations under the License.
+
+cd /d "@_bazel_src_root@"
+@IREE_BAZEL_EXECUTABLE@ @_bazel_startup_options@ %* || exit /b
diff --git a/packaging/python/dummy_exclude_from_package.py b/build_tools/cmake/bazel.sh.in
similarity index 84%
copy from packaging/python/dummy_exclude_from_package.py
copy to build_tools/cmake/bazel.sh.in
index 0ca47f8..ed2466f 100644
--- a/packaging/python/dummy_exclude_from_package.py
+++ b/build_tools/cmake/bazel.sh.in
@@ -1,3 +1,4 @@
+#!/bin/bash
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,3 +12,6 @@
# 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.
+
+cd "@_bazel_src_root@"
+exec '@IREE_BAZEL_EXECUTABLE@' @_bazel_startup_options@ "$@"
diff --git a/build_tools/cmake/configure_bazel.cmake b/build_tools/cmake/configure_bazel.cmake
new file mode 100644
index 0000000..528d244
--- /dev/null
+++ b/build_tools/cmake/configure_bazel.cmake
@@ -0,0 +1,158 @@
+# 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.
+
+set(IREE_BAZEL_EXECUTABLE "bazel"
+ CACHE STRING "Bazel executable to use for bazel builds")
+
+# iree_configure_bazel
+#
+# Configures the CMake binary directory to also contain a bazel build root.
+# The following files will be created:
+# bazel (shell script): Shell wrapper to invoke bazel
+# bazel.bat: Windows batch file to invoke bazel
+# bazelrc: The bazelrc to use for the build
+# bazel-out/: Bazel output directory
+# bazel-bin/: Symlink to the bin directory appropriate for the build mode
+#
+# Variables will be set in the parent scope:
+# IREE_BAZEL_WRAPPER: Executable wrapper to invoke to run bazel
+# IREE_BAZEL_BIN: Path to the bazel-bin directory
+function(iree_configure_bazel)
+ set(_bazel_output_base "${CMAKE_BINARY_DIR}/bazel-out")
+ set(_bazel_src_root "${CMAKE_SOURCE_DIR}")
+
+ # Use the utility to emit _bazelrc_file configuration options.
+ set(_bazelrc_file "${CMAKE_BINARY_DIR}/bazelrc")
+ execute_process(
+ RESULT_VARIABLE RC
+ COMMAND
+ "${Python3_EXECUTABLE}"
+ "${_bazel_src_root}/configure_bazel.py"
+ "${_bazelrc_file}"
+ )
+ if(NOT RC EQUAL 0)
+ message(FATAL_ERROR "Error running ${_bazel_src_root}/configure_bazel.py script")
+ endif()
+
+ # Now add an import to the configured.bazelrc to load the project-wide
+ # bazelrc file.
+ file(APPEND "${_bazelrc_file}" "
+import ${_bazel_src_root}/build_tools/bazel/iree.bazelrc
+")
+
+ # Note that we do allow a .bazelrc in the home directory (otherwise we
+ # would have --nohome_rc). This is mainly about disabling interference from
+ # interactive builds in the workspace.
+ set(_bazel_startup_options "--nosystem_rc --noworkspace_rc '--bazelrc=${_bazelrc_file}' '--output_base=${_bazel_output_base}'")
+
+ # And emit scripts to delegate to bazel.
+ set(IREE_BAZEL_WRAPPER "${CMAKE_BINARY_DIR}/bazel")
+ configure_file(
+ "${CMAKE_CURRENT_SOURCE_DIR}/build_tools/cmake/bazel.sh.in"
+ "${IREE_BAZEL_WRAPPER}"
+ )
+ configure_file(
+ "${CMAKE_CURRENT_SOURCE_DIR}/build_tools/cmake/bazel.bat.in"
+ "${IREE_BAZEL_WRAPPER}.bat"
+ )
+ if(NOT WIN32)
+ execute_process(
+ COMMAND chmod a+x "${IREE_BAZEL_WRAPPER}"
+ )
+ endif()
+
+ # Now ready to start bazel and ask it things.
+ message(STATUS "Detecting bazel version...")
+ execute_process(
+ RESULT_VARIABLE RC
+ OUTPUT_VARIABLE BAZEL_RELEASE
+ OUTPUT_STRIP_TRAILING_WHITESPACE
+ COMMAND
+ "${IREE_BAZEL_WRAPPER}" info release
+ )
+ if(NOT RC EQUAL 0)
+ message(FATAL_ERROR "Failed to launch bazel using wrapper ${IREE_BAZEL_WRAPPER}. Inspect that script and ensure bazel is installed properly.")
+ endif()
+ execute_process(
+ RESULT_VARIABLE RC
+ OUTPUT_VARIABLE IREE_BAZEL_BIN
+ OUTPUT_STRIP_TRAILING_WHITESPACE
+ COMMAND
+ "${IREE_BAZEL_WRAPPER}" info bazel-bin
+ )
+ if(NOT RC EQUAL 0)
+ message(FATAL_ERROR "Failed to run 'info bazel-bin' via ${IREE_BAZEL_WRAPPER}. Inspect that script and ensure bazel is installed properly.")
+ endif()
+ message(STATUS "Found bazel ${BAZEL_RELEASE}, bin directory: ${IREE_BAZEL_BIN}")
+ message(STATUS "Bazel wrapper script generated at: ${IREE_BAZEL_WRAPPER}")
+
+ # Build automation will use the IREE_BAZEL_BIN variable, but also drop a
+ # convenience symlink, since that is what people expect.
+ # And since bazel isn't nice enough to create it...
+ if(NOT WIN32)
+ execute_process(
+ RESULT_VARIABLE RC
+ COMMAND
+ ln -sf "${IREE_BAZEL_BIN}" "${CMAKE_CURRENT_BINARY_DIR}/bazel-bin"
+ )
+ if(NOT RC EQUAL 0)
+ message(WARNING "Failed to create convenience bazel-bin symlink")
+ endif()
+ endif()
+
+ set(IREE_BAZEL_WRAPPER "${IREE_BAZEL_WRAPPER}" PARENT_SCOPE)
+ set(IREE_BAZEL_BIN "${IREE_BAZEL_BIN}" PARENT_SCOPE)
+endfunction()
+
+# iree_add_bazel_invocation
+#
+# Adds a target to perform a bazel invocation, building a list of targets
+# and exporting pseudo targets for some results of the build.
+#
+# Parameters:
+# INVOCATION_TARGET: The target name for the custom invocation target.
+# BAZEL_TARGETS: List of bazel targets to build.
+# EXECUTABLE_PATHS: Paths under bazel-bin for executables. An equivalent
+# CMake imported executable target will be created for each by replacing
+# the "/" with "_".
+function(iree_add_bazel_invocation)
+ cmake_parse_arguments(ARG
+ ""
+ "INVOCATION_TARGET"
+ "BAZEL_TARGETS;EXECUTABLE_PATHS"
+ ${ARGN}
+ )
+
+ add_custom_target(${ARG_INVOCATION_TARGET}
+ USES_TERMINAL
+ COMMAND ${CMAKE_COMMAND} -E echo
+ "Starting bazel build of targets ${ARG_BAZEL_TARGETS}"
+ COMMAND "${IREE_BAZEL_WRAPPER}" build ${ARG_BAZEL_TARGETS}
+ COMMAND ${CMAKE_COMMAND} -E echo "Bazel build complete."
+ )
+
+ # Create an imported executable target for each binary path.
+ # Since the bazel directory namespace lines up with the cmake namespace,
+ # generate a cmake target name for each.
+ foreach(_executable_path ${ARG_EXECUTABLE_PATHS})
+ string(REPLACE "/" "_" _executable_target "${_executable_path}")
+ message(STATUS "Add bazel executable target ${_executable_target}")
+ add_executable(${_executable_target} IMPORTED GLOBAL)
+ set_target_properties(${_executable_target}
+ PROPERTIES IMPORTED_LOCATION
+ "${IREE_BAZEL_BIN}/${_executable_path}${CMAKE_EXECUTABLE_SUFFIX}"
+ )
+ add_dependencies(${_executable_target} ${ARG_INVOCATION_TARGET})
+ endforeach()
+endfunction()
diff --git a/build_tools/cmake/iree_cross_compile.cmake b/build_tools/cmake/iree_cross_compile.cmake
index f749f90..581f3d2 100644
--- a/build_tools/cmake/iree_cross_compile.cmake
+++ b/build_tools/cmake/iree_cross_compile.cmake
@@ -87,6 +87,7 @@
iree_to_bool(_CONFIG_BUILD_PYTHON_BINDINGS "${IREE_${CONFIG_NAME}_BUILD_PYTHON_BINDINGS}")
iree_to_bool(_CONFIG_BUILD_JAVA_BINDINGS "${IREE_${CONFIG_NAME}_BUILD_JAVA_BINDINGS}")
iree_to_bool(_CONFIG_BUILD_EXPERIMENTAL "${IREE_${CONFIG_NAME}_BUILD_EXPERIMENTAL}")
+ iree_to_bool(_CONFIG_BUILD_TENSORFLOW_COMPILER "${IREE_${CONFIG_NAME}_BUILD_TENSORFLOW_COMPILER}")
# Escape semicolons in the targets list so that CMake doesn't expand them to
# spaces.
@@ -119,6 +120,7 @@
-DIREE_BUILD_PYTHON_BINDINGS=${_CONFIG_BUILD_PYTHON_BINDINGS}
-DIREE_BUILD_JAVA_BINDINGS=${_CONFIG_BUILD_JAVA_BINDINGS}
-DIREE_BUILD_EXPERIMENTAL=${_CONFIG_BUILD_EXPERIMENTAL}
+ -DIREE_BUILD_TENSORFLOW_COMPILER=${_CONFIG_BUILD_TENSORFLOW_COMPILER}
# LINT.ThenChange(
# https://github.com/google/iree/tree/main/CMakeLists.txt:iree_options,
# https://github.com/google/iree/tree/main/build_tools/cmake/iree_cross_compile.cmake:iree_cross_compile_options,
diff --git a/build_tools/cmake/iree_macros.cmake b/build_tools/cmake/iree_macros.cmake
index 04466b5..8f4e5d2 100644
--- a/build_tools/cmake/iree_macros.cmake
+++ b/build_tools/cmake/iree_macros.cmake
@@ -265,6 +265,49 @@
endfunction()
#-------------------------------------------------------------------------------
+# Tool symlinks
+#-------------------------------------------------------------------------------
+
+# iree_symlink_tool
+#
+# Adds a command to TARGET which symlinks a tool from elsewhere
+# (FROM_TOOL_TARGET_NAME) to a local file name (TO_EXE_NAME) in the current
+# binary directory.
+#
+# Parameters:
+# TARGET: Local target to which to add the symlink command (i.e. an
+# iree_py_library, etc).
+# FROM_TOOL_TARGET: Target of the tool executable that is the source of the
+# link.
+# TO_EXE_NAME: The executable name to output in the current binary dir.
+function(iree_symlink_tool)
+ cmake_parse_arguments(
+ ARG
+ ""
+ "TARGET;FROM_TOOL_TARGET;TO_EXE_NAME"
+ ""
+ ${ARGN}
+ )
+
+ # Transform TARGET
+ iree_package_ns(_PACKAGE_NS)
+ iree_package_name(_PACKAGE_NAME)
+ set(_TARGET "${_PACKAGE_NAME}_${ARG_TARGET}")
+ set(_FROM_TOOL_TARGET ${ARG_FROM_TOOL_TARGET})
+
+ add_custom_command(
+ TARGET "${_TARGET}"
+ BYPRODUCTS
+ "${CMAKE_CURRENT_BINARY_DIR}/${ARG_TO_EXE_NAME}${CMAKE_EXECUTABLE_SUFFIX}"
+ COMMAND
+ ${CMAKE_COMMAND} -E create_symlink
+ "$<TARGET_FILE:${_FROM_TOOL_TARGET}>"
+ "${CMAKE_CURRENT_BINARY_DIR}/${ARG_TO_EXE_NAME}${CMAKE_EXECUTABLE_SUFFIX}"
+ )
+endfunction()
+
+
+#-------------------------------------------------------------------------------
# Tests
#-------------------------------------------------------------------------------
diff --git a/build_tools/cmake/iree_multipy.cmake b/build_tools/cmake/iree_multipy.cmake
index 1ac1028..3456fad 100644
--- a/build_tools/cmake/iree_multipy.cmake
+++ b/build_tools/cmake/iree_multipy.cmake
@@ -23,12 +23,20 @@
# Note that this is using the pybind11 configuration vars, which creates
# a fragile dependency. It would be better to derive these locally.
if(Python3_FOUND)
- set(IREE_MULTIPY_DEFAULT_EXECUTABLE "${PYTHON_EXECUTABLE}" CACHE INTERNAL "Python executable" )
- set(IREE_MULTIPY_DEFAULT_INCLUDE_DIRS "${PYTHON_INCLUDE_DIRS}" CACHE INTERNAL "Python include dirs" )
- set(IREE_MULTIPY_DEFAULT_LIBRARIES "${PYTHON_LIBRARIES}" CACHE INTERNAL "Python libraries")
- set(IREE_MULTIPY_DEFAULT_PREFIX "${PYTHON_MODULE_PREFIX}" CACHE INTERNAL "Python module prefix")
- set(IREE_MULTIPY_DEFAULT_SUFFIX "${PYTHON_MODULE_SUFFIX}" CACHE INTERNAL "Python module suffix")
- set(IREE_MULTIPY_DEFAULT_EXTENSION "${PYTHON_MODULE_EXTENSION}" CACHE INTERNAL "Python module extension")
+ set(IREE_MULTIPY_DEFAULT_EXECUTABLE "${Python3_EXECUTABLE}" CACHE INTERNAL "Python executable" )
+ set(IREE_MULTIPY_DEFAULT_INCLUDE_DIRS "${Python3_INCLUDE_DIRS}" CACHE INTERNAL "Python include dirs" )
+ set(IREE_MULTIPY_DEFAULT_LIBRARIES "${Python3_LIBRARIES}" CACHE INTERNAL "Python libraries")
+ set(IREE_MULTIPY_DEFAULT_PREFIX "${Python3_MODULE_PREFIX}" CACHE INTERNAL "Python module prefix")
+ set(IREE_MULTIPY_DEFAULT_SUFFIX "${Python3_MODULE_SUFFIX}" CACHE INTERNAL "Python module suffix")
+ # CMake 3.19 and there-abouts does define Python3_SOABI, but get it
+ # ourselves for compatibility.
+ execute_process(
+ OUTPUT_VARIABLE _FOUND_DEFAULT_EXTENSION
+ OUTPUT_STRIP_TRAILING_WHITESPACE
+ COMMAND
+ "${Python3_EXECUTABLE}" -c "import sysconfig;print(sysconfig.get_config_var('EXT_SUFFIX'))"
+ )
+ set(IREE_MULTIPY_DEFAULT_EXTENSION "${_FOUND_DEFAULT_EXTENSION}" CACHE INTERNAL "Python module extension")
endif()
if(IREE_MULTIPY_VERSIONS)
@@ -51,13 +59,13 @@
# Check for required settings.
if(NOT IREE_MULTIPY_${V}_INCLUDE_DIRS)
- message(FATAL " MULTIPY version ${V}: No IREE_MULTIPY_${VER}_EXECUTABLE var")
+ message(FATAL_ERROR " MULTIPY version ${V}: No IREE_MULTIPY_${VER}_EXECUTABLE var")
endif()
if(NOT IREE_MULTIPY_${V}_INCLUDE_DIRS)
- message(FATAL " MULTIPY version ${V}: No IREE_MULTIPY_${VER}_INCLUDE_DIRS var")
+ message(FATAL_ERROR " MULTIPY version ${V}: No IREE_MULTIPY_${VER}_INCLUDE_DIRS var")
endif()
if(NOT IREE_MULTIPY_${V}_EXTENSION)
- message(FATAL " MULTIPY version ${V}: No IREE_MULTIPY_${VER}_EXTENSION var")
+ message(FATAL_ERROR " MULTIPY version ${V}: No IREE_MULTIPY_${VER}_EXTENSION var")
endif()
endforeach()
endfunction()
@@ -237,14 +245,20 @@
iree_package_name(_PACKAGE_NAME)
set(_NAME "${_PACKAGE_NAME}_${ARG_NAME}")
- # Add path to each source file
- list(TRANSFORM ARG_SRCS PREPEND "${CMAKE_CURRENT_SOURCE_DIR}/")
-
add_custom_target(${_NAME} ALL
- COMMAND ${CMAKE_COMMAND} -E copy ${ARG_SRCS} "${CMAKE_CURRENT_BINARY_DIR}/"
DEPENDS ${ARG_DEPS}
)
+ # Symlink each file as its own target.
+ foreach(SRC_FILE ${ARG_SRCS})
+ add_custom_command(
+ TARGET ${_NAME}
+ COMMAND ${CMAKE_COMMAND} -E create_symlink
+ "${CMAKE_CURRENT_SOURCE_DIR}/${SRC_FILE}" "${CMAKE_CURRENT_BINARY_DIR}/${SRC_FILE}"
+ BYPRODUCTS "${CMAKE_CURRENT_BINARY_DIR}/${SRC_FILE}"
+ )
+ endforeach()
+
# Add PYEXT_DEPS.
if(${ARG_PYEXT_DEPS})
foreach(V ${IREE_MULTIPY_VERSIONS_EFFECTIVE})
diff --git a/configure_bazel.py b/configure_bazel.py
index f0bac93..18be4a4 100644
--- a/configure_bazel.py
+++ b/configure_bazel.py
@@ -49,7 +49,7 @@
user_site = subprocess.check_output(
[sys.executable, "-m", "site", "--user-site"]).decode("utf-8").strip()
print("Found user site directory:", user_site)
- except OSError:
+ except subprocess.CalledProcessError:
print("Could not resolve user site directory")
return
print("build --action_env PYTHONPATH=\"{}\"".format(
@@ -57,7 +57,10 @@
file=bazelrc)
-local_bazelrc = os.path.join(os.path.dirname(__file__), "configured.bazelrc")
+if len(sys.argv) > 1:
+ local_bazelrc = sys.argv[1]
+else:
+ local_bazelrc = os.path.join(os.path.dirname(__file__), "configured.bazelrc")
with open(local_bazelrc, "wt") as bazelrc:
write_platform(bazelrc)
write_python_bin(bazelrc)
diff --git a/docs/get_started/cmake_options_and_variables.md b/docs/get_started/cmake_options_and_variables.md
index 18fff58..ad523a8 100644
--- a/docs/get_started/cmake_options_and_variables.md
+++ b/docs/get_started/cmake_options_and_variables.md
@@ -105,6 +105,15 @@
`MLIR_DIR` needs to be passed and that LLVM needs to be compiled with
`LLVM_ENABLE_RTTI` set to `ON`.
+#### `IREE_BUILD_TENSORFLOW_COMPILER`:BOOL
+
+Enables building of the TensorFlow to IREE compiler under
+`integrations/tensorflow`, including some native binaries and Python packages.
+Note that TensorFlow's build system is bazel-based and this will require a
+functioning `bazel` installation. A `bazel` wrapper script will be emitted
+in your build directory and `bazel-bin` link will point to artifacts. This
+can be used to manually invoke additional bazel actions if desired.
+
## MLIR-specific CMake Options and Variables
#### `MLIR_DIR`:STRING
diff --git a/docs/get_started/getting_started_python.md b/docs/get_started/getting_started_python.md
index af8a25b..955a179 100644
--- a/docs/get_started/getting_started_python.md
+++ b/docs/get_started/getting_started_python.md
@@ -1,11 +1,14 @@
# Getting Started with Python
-IREE has Python bindings geared towards lower level compiler interop that are
-not intended to be a public API, and integration with Python frontends such as
-TensorFlow.
+ NOTE: Iree's Python API is currently being reworked. Some of these
+ instructions may be in a state of flux as they document the end state.
-We do not yet provide a pip package for easy installation, so to use IREE's
-Python bindings you must build from source.
+IREE has two primary Python APIs:
+
+* Compiler API: `pyiree.compiler2`, `pyiree.compiler2.tf`
+* Runtime API: `pyiree.tf`
+
+There are additional ancillary modules that are not part of the public API.
## Prerequisites
@@ -13,22 +16,41 @@
[getting started guides](../get-started) for instructions.
> Note:<br>
-> Support is best with Bazel.
-> For CMake (excluding TensorFlow), set the `IREE_BUILD_PYTHON_BINDINGS` option.
+> Support is only complete with CMake.
+
+Minimally, the following CMake flags must be specified:
+
+* `-DIREE_BUILD_PYTHON_BINDINGS=ON`
+* `-DIREE_BUILD_TENSORFLOW_COMPILER=ON` : Optional. Also builds the
+ TensorFlow compiler integration.
+
+If building any parts of TensorFlow, you must have a working `bazel` command
+on your path. See the `.bazelversion` file at the root of the project for the
+recommended version.
## Python Setup
Install a recent version of [Python 3](https://www.python.org/downloads/) and
[pip](https://pip.pypa.io/en/stable/installing/), if needed.
+(Recommended) Setup a virtual environment (use your preferred mechanism):
+
+```shell
+# Note that venv is only available in python3 and is therefore a good check
+# that you are in fact running a python3 binary.
+python -m venv .venv
+source .venv/bin/activate
+# When done: run 'deactivate'
+```
+
Install packages:
```shell
-$ python3 -m pip install --upgrade pip
-$ python3 -m pip install numpy
+$ python -m pip install --upgrade pip
+$ python -m pip install numpy absl-py
# If using the TensorFlow integration
-$ python3 -m pip install tf-nightly
+$ python -m pip install tf-nightly
```
## Running Python Tests
@@ -40,29 +62,60 @@
$ ctest -L bindings/python
```
-To run tests for core Python bindings built with Bazel:
-
-```shell
-$ bazel test bindings/python/...
-```
-
To run tests for the TensorFlow integration, which include end-to-end backend
comparison tests:
```shell
-# Exclude tests that are skipped in the Kokoro CI
-$ bazel test \
- --build_tag_filters="-nokokoro" \
- --test_tag_filters="-nokokoro" \
- --define=iree_tensorflow=true \
- integrations/tensorflow/...
+cd build
+# TODO: Revisit once more patches land.
+ctest -L integrations/tensorflow/e2e
+
+# Or run individually as:
+export PYTHONPATH=bindings/python # In build dir
+python integrations/tensorflow/e2e/simple_arithmetic_test.py \
+ --target_backends=iree_vmla --artifacts_dir=/tmp/artifacts
```
+
## Using Colab
-See
-[start_colab_kernel.py](https://github.com/google/iree/blob/main/colab/start_colab_kernel.py)
-and [Using Colab](../using_iree/using_colab.md) for setup instructions, then
-take a look through the
-[Colab directory](https://github.com/google/iree/tree/main/colab) for some
-sample notebooks.
+There are some sample colabs in the `colab` folder. If you have built the
+project with CMake/ninja and set your `PYTHONPATH` to the `bindings/python`
+directory in the build dir (or installed per below), you should be able to
+start a kernel by following the stock instructions at
+https://colab.research.google.com/ .
+
+
+## Installing and Packaging
+
+There is a `setup.py` in the `bindings/python` directory under the build dir.
+To install into your (hopefully isolated) virtual env:
+
+```shell
+python bindings/python/setup.py install
+```
+
+To create wheels (platform dependent and locked to your Python version
+without further config):
+
+```shell
+python bindings/python/setup.py bdist_wheel
+```
+
+Note that it is often helpful to differentiate between the environment used to
+build and the one used to install. While this is just "normal" python
+knowledge, here is an incantation to do so:
+
+```shell
+# From parent/build environment.
+python -m pip freeze > /tmp/requirements.txt
+deactivate # If already in an environment
+
+# Enter new scratch environment.
+python -m venv ./.venv-scratch
+source ./.venv-scratch/bin/activate
+python -m pip install -r /tmp/requirements.txt
+
+# Install IREE into the new environment.
+python bindings/python/setup.py install
+```
diff --git a/experimental/ModelBuilder/test/CMakeLists.txt b/experimental/ModelBuilder/test/CMakeLists.txt
index 09ab2f2..7af3c3b 100644
--- a/experimental/ModelBuilder/test/CMakeLists.txt
+++ b/experimental/ModelBuilder/test/CMakeLists.txt
@@ -43,7 +43,6 @@
"TestDotProdJIT.cpp"
DEPS
LLVMSupport
- MLIRAllDialects
MLIREDSC
MLIRIR
MLIRSCFTransforms
@@ -59,9 +58,9 @@
SRCS
"TestVectorTransfersJIT.cpp"
DEPS
- runtime-support.so
LLVMSupport
- MLIRAllDialects
+ # TODO(thomasraoux): Fix dependecy to shared library.
+ # runtime-support.so
MLIREDSC
MLIRIR
MLIRSCFTransforms
@@ -77,11 +76,9 @@
SRCS
"TestMNISTJIT.cpp"
DEPS
- MLIRAllDialects
MLIREDSC
MLIRIR
MLIRSCFTransforms
- MLIRmlir_runner_utils
experimental::ModelBuilder
experimental::ModelBuilder::ModelRunner
)
@@ -95,7 +92,6 @@
"TestSimpleJIT.cpp"
DEPS
LLVMSupport
- MLIRAllDialects
MLIREDSC
MLIRIR
MLIRSCFTransforms
@@ -112,17 +108,11 @@
"TestSimpleJITVulkan.cpp"
DEPS
LLVMSupport
- MLIRAllDialects
MLIRIR
MLIRParser
MLIRSPIRV
- MLIRmlir_runner_utils
experimental::ModelBuilder
experimental::ModelBuilder::ModelRunner
- iree::base::initializer
- iree::hal::llvmjit::llvmjit_driver_module
- iree::hal::vmla::vmla_driver_module
- iree::hal::vulkan::vulkan_driver_module
vulkan-runtime-wrappers
)
@@ -135,7 +125,6 @@
"TestMatMulVulkan.cpp"
DEPS
LLVMSupport
- MLIRAllDialects
MLIRExecutionEngine
MLIRGPU
MLIRGPUToSPIRVTransforms
@@ -153,15 +142,11 @@
MLIRTargetLLVMIR
MLIRTransformUtils
MLIRVectorToLLVM
- MLIRmlir_runner_utils
experimental::ModelBuilder
experimental::ModelBuilder::ModelRunner
experimental::ModelBuilder::VulkanLaunchWrapper
- iree::base::initializer
iree::compiler::Conversion::LinalgToSPIRV
- iree::hal::llvmjit::llvmjit_driver_module
- iree::hal::vmla::vmla_driver_module
- iree::hal::vulkan::vulkan_driver_module
+ iree::tools::init_mlir_passes_and_dialects
vulkan-runtime-wrappers
)
@@ -174,7 +159,6 @@
"TestVectorToGPU.cpp"
DEPS
LLVMSupport
- MLIRAllDialects
MLIRExecutionEngine
MLIRGPU
MLIRGPUToVulkanTransforms
@@ -189,16 +173,12 @@
MLIRStandardToSPIRVTransforms
MLIRTransformUtils
MLIRVector
- MLIRmlir_runner_utils
experimental::ModelBuilder
experimental::ModelBuilder::ModelRunner
experimental::ModelBuilder::VulkanLaunchWrapper
- iree::base::initializer
iree::compiler::Conversion::CodegenUtils
iree::compiler::Conversion::LinalgToSPIRV
- iree::hal::llvmjit::llvmjit_driver_module
- iree::hal::vmla::vmla_driver_module
- iree::hal::vulkan::vulkan_driver_module
+ iree::tools::init_mlir_passes_and_dialects
vulkan-runtime-wrappers
)
@@ -211,7 +191,6 @@
"BenchMatMulVectorGPU.cpp"
DEPS
LLVMSupport
- MLIRAllDialects
MLIRExecutionEngine
MLIRGPU
MLIRGPUToVulkanTransforms
@@ -226,16 +205,12 @@
MLIRStandardToSPIRVTransforms
MLIRTransformUtils
MLIRVector
- MLIRmlir_runner_utils
experimental::ModelBuilder
experimental::ModelBuilder::ModelRunner
experimental::ModelBuilder::VulkanLaunchWrapper
- iree::base::initializer
iree::compiler::Conversion::CodegenUtils
iree::compiler::Conversion::LinalgToSPIRV
- iree::hal::llvmjit::llvmjit_driver_module
- iree::hal::vmla::vmla_driver_module
- iree::hal::vulkan::vulkan_driver_module
+ iree::tools::init_mlir_passes_and_dialects
vulkan-runtime-wrappers
)
@@ -247,7 +222,6 @@
SRCS
"TestSimpleMLIR.cpp"
DEPS
- MLIRAllDialects
experimental::ModelBuilder
experimental::ModelBuilder::ModelRunner
)
@@ -260,7 +234,6 @@
SRCS
"BenchMatVecVectorJIT.cpp"
DEPS
- MLIRAllDialects
MLIREDSC
MLIRIR
benchmark
@@ -276,7 +249,6 @@
SRCS
"BenchMatMulVectorJIT.cpp"
DEPS
- MLIRAllDialects
MLIREDSC
MLIRIR
benchmark
@@ -292,7 +264,6 @@
SRCS
"BenchMatMulVectorColumnMajorLLVMIntrinsicsJIT.cpp"
DEPS
- MLIRAllDialects
MLIREDSC
MLIRIR
benchmark
diff --git a/integrations/tensorflow/CMakeLists.txt b/integrations/tensorflow/CMakeLists.txt
new file mode 100644
index 0000000..6194ee3
--- /dev/null
+++ b/integrations/tensorflow/CMakeLists.txt
@@ -0,0 +1,36 @@
+# 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.
+
+# TensorFlow builds through bazel, and IREE maintains all of its TensorFlow
+# dependent code under this directory tree. The CMake support is limited to
+# compiler binaries and python bindings.
+#
+# Bazel is a beast that likes to be the center of the universe. There is some
+# fragility in delegating to it in this fashion.
+#
+# If this directory is included, then building TensorFlow is assumed (the
+# config option happens at the higher level).
+
+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
+)
+
+if(${IREE_BUILD_PYTHON_BINDINGS})
+ add_subdirectory(bindings/python)
+endif()
diff --git a/bindings/python/pyiree/CMakeLists.txt b/integrations/tensorflow/bindings/python/CMakeLists.txt
similarity index 62%
copy from bindings/python/pyiree/CMakeLists.txt
copy to integrations/tensorflow/bindings/python/CMakeLists.txt
index 55ac43d..31b61a2 100644
--- a/bindings/python/pyiree/CMakeLists.txt
+++ b/integrations/tensorflow/bindings/python/CMakeLists.txt
@@ -12,9 +12,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-add_subdirectory(common)
-add_subdirectory(rt)
+# Overlays a subdirectory into the main python bindings directory.
+function(_add_overlay_subdirectory dir)
+ # Overlay binary directories onto the main bindings directory.
+ set(_MAIN_PYTHON_DIR "${CMAKE_BINARY_DIR}/bindings/python")
+ add_subdirectory(${dir} "${_MAIN_PYTHON_DIR}/${dir}")
+endfunction()
-if(${IREE_BUILD_COMPILER})
- add_subdirectory(compiler)
-endif()
+_add_overlay_subdirectory(pyiree/tools/tf)
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/tf/.skip_bazel_to_cmake b/integrations/tensorflow/bindings/python/pyiree/tools/tf/.skip_bazel_to_cmake
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/integrations/tensorflow/bindings/python/pyiree/tools/tf/.skip_bazel_to_cmake
diff --git a/bindings/python/pyiree/CMakeLists.txt b/integrations/tensorflow/bindings/python/pyiree/tools/tf/CMakeLists.txt
similarity index 70%
copy from bindings/python/pyiree/CMakeLists.txt
copy to integrations/tensorflow/bindings/python/pyiree/tools/tf/CMakeLists.txt
index 55ac43d..a10869b 100644
--- a/bindings/python/pyiree/CMakeLists.txt
+++ b/integrations/tensorflow/bindings/python/pyiree/tools/tf/CMakeLists.txt
@@ -12,9 +12,17 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-add_subdirectory(common)
-add_subdirectory(rt)
+iree_py_library(
+ NAME
+ tf
+ SRCS
+ "__init__.py"
+ DEPS
+ integrations_iree_tensorflow_importers
+)
-if(${IREE_BUILD_COMPILER})
- add_subdirectory(compiler)
-endif()
+iree_symlink_tool(
+ TARGET tf
+ FROM_TOOL_TARGET integrations_tensorflow_compiler_iree-tf-import
+ TO_EXE_NAME iree-tf-import
+)
diff --git a/bindings/python/pyiree/CMakeLists.txt b/integrations/tensorflow/bindings/python/pyiree/tools/tf/__init__.py
similarity index 63%
copy from bindings/python/pyiree/CMakeLists.txt
copy to integrations/tensorflow/bindings/python/pyiree/tools/tf/__init__.py
index 55ac43d..ceb37ea 100644
--- a/bindings/python/pyiree/CMakeLists.txt
+++ b/integrations/tensorflow/bindings/python/pyiree/tools/tf/__init__.py
@@ -1,3 +1,6 @@
+# Lint-as: python3
+"""TensorFlow tools."""
+
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,9 +15,15 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-add_subdirectory(common)
-add_subdirectory(rt)
+from typing import Optional
-if(${IREE_BUILD_COMPILER})
- add_subdirectory(compiler)
-endif()
+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/iree/compiler/Conversion/Common/VectorTransferOptimization.cpp b/iree/compiler/Conversion/Common/VectorTransferOptimization.cpp
index 91ef5a0..85aa505 100644
--- a/iree/compiler/Conversion/Common/VectorTransferOptimization.cpp
+++ b/iree/compiler/Conversion/Common/VectorTransferOptimization.cpp
@@ -12,17 +12,57 @@
// See the License for the specific language governing permissions and
// limitations under the License.
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
#include "mlir/Dialect/Vector/VectorTransforms.h"
#include "mlir/Pass/Pass.h"
namespace mlir {
namespace iree_compiler {
+// Return true if all the uses of op are either Store/transfer_write.
+// There can be SubviewOp users as long as all its users are also
+// StoreOp/transfer_write. If return true it also fills out the uses, if it
+// returns false uses is unchanged.
+static bool allUsesAreStores(Operation* op, std::vector<Operation*>& uses) {
+ std::vector<Operation*> opUses;
+ for (OpOperand& use : op->getUses()) {
+ Operation* useOp = use.getOwner();
+ if (isa<vector::TransferWriteOp, StoreOp>(useOp) ||
+ (isa<SubViewOp>(useOp) && allUsesAreStores(useOp, opUses))) {
+ opUses.push_back(useOp);
+ continue;
+ }
+ return false;
+ }
+ uses.insert(uses.end(), opUses.begin(), opUses.end());
+ return true;
+}
+
+// Track temporary allocations that are never read from. If this is the case
+// it means both the allocations and associated stores can be removed.
+static void eraseDeadAllocAndStores(FuncOp funcOp) {
+ std::vector<Operation*> opToErase;
+ funcOp.walk([&](AllocOp op) {
+ if (allUsesAreStores(op, opToErase)) {
+ opToErase.push_back(op.getOperation());
+ }
+ });
+ for (Operation* op : opToErase) {
+ op->erase();
+ }
+}
+
namespace {
struct VectorTransferOptimizationPass
: public PassWrapper<VectorTransferOptimizationPass, FunctionPass> {
- void runOnFunction() override { vector::transferOpflowOpt(getFunction()); }
+ void runOnFunction() override {
+ vector::transferOpflowOpt(getFunction());
+ // Delete potential dead alloc and associated ops after store to load
+ // forwarding.
+ eraseDeadAllocAndStores(getFunction());
+ }
};
} // namespace
diff --git a/iree/compiler/Conversion/HLOToLinalg/BUILD b/iree/compiler/Conversion/HLOToLinalg/BUILD
index a1f9de8..56a034d 100644
--- a/iree/compiler/Conversion/HLOToLinalg/BUILD
+++ b/iree/compiler/Conversion/HLOToLinalg/BUILD
@@ -19,11 +19,35 @@
)
cc_library(
- name = "HLOToLinalg",
+ name = "HLOToLinalgOnTensors",
srcs = [
"FusionOfTensorOps.cpp",
- "HLOToLinalgOnBuffers.cpp",
"HLOToLinalgOnTensors.cpp",
+ ],
+ hdrs = [
+ "HLOToLinalgOnTensorPasses.h",
+ ],
+ deps = [
+ "//iree/compiler/Dialect/HAL/IR",
+ "//iree/compiler/Dialect/HAL/IR:HALDialect",
+ "@llvm-project//mlir:Affine",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:LinalgOps",
+ "@llvm-project//mlir:LinalgTransforms",
+ "@llvm-project//mlir:Pass",
+ "@llvm-project//mlir:StandardOps",
+ "@llvm-project//mlir:Support",
+ "@llvm-project//mlir:Transforms",
+ "@org_tensorflow//tensorflow/compiler/mlir/hlo",
+ "@org_tensorflow//tensorflow/compiler/mlir/hlo:legalize_to_linalg",
+ "@org_tensorflow//tensorflow/compiler/mlir/hlo:map_lmhlo_to_scalar_op",
+ ],
+)
+
+cc_library(
+ name = "HLOToLinalg",
+ srcs = [
+ "HLOToLinalgOnBuffers.cpp",
"Passes.cpp",
"ResolveShapeOps.cpp",
],
@@ -31,6 +55,7 @@
"Passes.h",
],
deps = [
+ ":HLOToLinalgOnTensors",
"//iree/compiler/Conversion/CodegenUtils",
"//iree/compiler/Conversion/HLOToHLO",
"//iree/compiler/Dialect/HAL/IR",
@@ -38,6 +63,7 @@
"//iree/compiler/Dialect/IREE/IR",
"//iree/compiler/Dialect/Shape/IR",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:Affine",
"@llvm-project//mlir:DialectUtils",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:LinalgOps",
@@ -48,7 +74,6 @@
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
"@org_tensorflow//tensorflow/compiler/mlir/hlo",
- "@org_tensorflow//tensorflow/compiler/mlir/hlo:legalize_to_linalg",
"@org_tensorflow//tensorflow/compiler/mlir/hlo:map_lmhlo_to_scalar_op",
],
)
diff --git a/iree/compiler/Conversion/HLOToLinalg/CMakeLists.txt b/iree/compiler/Conversion/HLOToLinalg/CMakeLists.txt
index 5f24d65..8984a07 100644
--- a/iree/compiler/Conversion/HLOToLinalg/CMakeLists.txt
+++ b/iree/compiler/Conversion/HLOToLinalg/CMakeLists.txt
@@ -16,17 +16,40 @@
iree_cc_library(
NAME
+ HLOToLinalgOnTensors
+ HDRS
+ "HLOToLinalgOnTensorPasses.h"
+ SRCS
+ "FusionOfTensorOps.cpp"
+ "HLOToLinalgOnTensors.cpp"
+ DEPS
+ MLIRAffine
+ MLIRIR
+ MLIRLinalg
+ MLIRLinalgTransforms
+ MLIRPass
+ MLIRStandard
+ MLIRSupport
+ MLIRTransforms
+ iree::compiler::Dialect::HAL::IR
+ iree::compiler::Dialect::HAL::IR::HALDialect
+ tensorflow::mlir_hlo
+ PUBLIC
+)
+
+iree_cc_library(
+ NAME
HLOToLinalg
HDRS
"Passes.h"
SRCS
- "FusionOfTensorOps.cpp"
"HLOToLinalgOnBuffers.cpp"
- "HLOToLinalgOnTensors.cpp"
"Passes.cpp"
"ResolveShapeOps.cpp"
DEPS
+ ::HLOToLinalgOnTensors
LLVMSupport
+ MLIRAffine
MLIRIR
MLIRLinalg
MLIRLinalgTransforms
diff --git a/iree/compiler/Conversion/HLOToLinalg/FusionOfTensorOps.cpp b/iree/compiler/Conversion/HLOToLinalg/FusionOfTensorOps.cpp
index f2f7f3f..05a6c87 100644
--- a/iree/compiler/Conversion/HLOToLinalg/FusionOfTensorOps.cpp
+++ b/iree/compiler/Conversion/HLOToLinalg/FusionOfTensorOps.cpp
@@ -20,9 +20,10 @@
//
//===----------------------------------------------------------------------===//
-#include "iree/compiler/Conversion/HLOToLinalg/Passes.h"
+#include "iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensorPasses.h"
#include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -67,7 +68,8 @@
struct FusionOfTensorOpsPass
: public PassWrapper<FusionOfTensorOpsPass, OperationPass<>> {
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<linalg::LinalgDialect, IREE::HAL::HALDialect>();
+ registry
+ .insert<AffineDialect, IREE::HAL::HALDialect, linalg::LinalgDialect>();
}
void runOnOperation() override {
diff --git a/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensorPasses.h b/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensorPasses.h
new file mode 100644
index 0000000..5057bdc
--- /dev/null
+++ b/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensorPasses.h
@@ -0,0 +1,48 @@
+// 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.
+
+//===- HLOToLinalgOnTensorsPasses.h - Passes to convert from HLO To Linalg ===//
+//
+// IREE specific passes used in the HLO to Linalg on tensors conversion as well
+// as fusion.
+//
+//===----------------------------------------------------------------------===//
+#ifndef IREE_COMPILER_CONVERSION_HLOTOLINALGONTENSORS_PASSES_H_
+#define IREE_COMPILER_CONVERSION_HLOTOLINALGONTENSORS_PASSES_H_
+
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+/// Creates a pass to fuse operations on tensors.
+std::unique_ptr<Pass> createFusionOfTensorOpsPass();
+
+/// Creates XLA-HLO to Linalg on tensors transformation pass.
+std::unique_ptr<OperationPass<FuncOp>> createHLOToLinalgOnTensorsPass();
+
+/// Populates the patterns that convert from MHLO to Linalg on tensors. Imports
+/// patterns from XLA, as well as some IREE specific modifications.
+void populateHLOToLinalgOnTensorsConversionPatterns(
+ MLIRContext *context, OwningRewritePatternList &patterns);
+
+/// Populated passes to convert from MHLO to Linalg on tensors as well as fusion
+/// of the converted operations.
+void addHLOToLinalgOnTensorsPasses(OpPassManager &pm);
+
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_CONVERSION_HLOTOLINALGONTENSORS_PASSES_H_
diff --git a/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensors.cpp b/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensors.cpp
index a58a877..d86b947 100644
--- a/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensors.cpp
+++ b/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensors.cpp
@@ -21,7 +21,7 @@
//===----------------------------------------------------------------------===//
#include <memory>
-#include "iree/compiler/Conversion/HLOToLinalg/Passes.h"
+#include "iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensorPasses.h"
#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
#include "mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
diff --git a/iree/compiler/Conversion/HLOToLinalg/Passes.cpp b/iree/compiler/Conversion/HLOToLinalg/Passes.cpp
index 2bb1e30..ca0842b 100644
--- a/iree/compiler/Conversion/HLOToLinalg/Passes.cpp
+++ b/iree/compiler/Conversion/HLOToLinalg/Passes.cpp
@@ -15,6 +15,7 @@
#include "iree/compiler/Conversion/HLOToLinalg/Passes.h"
#include "iree/compiler/Conversion/HLOToHLO/Passes.h"
+#include "iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensorPasses.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/Passes.h"
@@ -23,12 +24,16 @@
namespace iree_compiler {
void addHLOToLinalgOnBuffersPasses(OpPassManager &pm) {
+ addHLOToLinalgOnTensorsPasses(pm);
+ pm.addNestedPass<FuncOp>(createHLOToLinalgOnBuffersPass());
+}
+
+void addHLOToLinalgOnTensorsPasses(OpPassManager &pm) {
pm.addNestedPass<FuncOp>(createHLOToLinalgOnTensorsPass());
pm.addNestedPass<FuncOp>(createLinalgFoldUnitExtentDimsPass());
pm.addNestedPass<FuncOp>(createCanonicalizerPass());
pm.addNestedPass<FuncOp>(createFusionOfTensorOpsPass());
pm.addNestedPass<FuncOp>(createCSEPass());
- pm.addNestedPass<FuncOp>(createHLOToLinalgOnBuffersPass());
}
static PassPipelineRegistration<> hloToLinalgOnBuffersPipeline(
diff --git a/iree/compiler/Conversion/HLOToLinalg/Passes.h b/iree/compiler/Conversion/HLOToLinalg/Passes.h
index 4c048b0..2119e42 100644
--- a/iree/compiler/Conversion/HLOToLinalg/Passes.h
+++ b/iree/compiler/Conversion/HLOToLinalg/Passes.h
@@ -27,24 +27,13 @@
namespace mlir {
namespace iree_compiler {
-/// Creates a pass to fuse operations on tensors.
-std::unique_ptr<Pass> createFusionOfTensorOpsPass();
-
/// Creates XLA-HLO to Linalg on buffers transformation pass.
std::unique_ptr<OperationPass<FuncOp>> createHLOToLinalgOnBuffersPass();
-/// Creates XLA-HLO to Linalg on tensors transformation pass.
-std::unique_ptr<OperationPass<FuncOp>> createHLOToLinalgOnTensorsPass();
-
/// Resolves shape related ops (std.dim, shapex.tie_shape, etc.) by tracing
/// them back to the original HAL interface bindings.
std::unique_ptr<OperationPass<FuncOp>> createResolveShapeOpsPass();
-/// Populates the patterns that convert from XLA to Linalg on tensors. Imports
-/// patterns from XLA, as well as some IREE specific modifications.
-void populateHLOToLinalgOnTensorsConversionPatterns(
- MLIRContext *context, OwningRewritePatternList &patterns);
-
/// Populates the patterns that convert from XLA to Linalg on buffers. Currently
/// only implements conversions when the XLA op is the only op XLA op in the
/// dispatch region.
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp b/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp
index 25d53b6..7caca66 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp
@@ -597,6 +597,16 @@
DISPATCH(vector::TransferWriteOp)
#undef DISPATCH
+
+ if (op->hasTrait<OpTrait::ElementwiseMappable>() &&
+ op->getNumResults() == 1) {
+ if (auto vecType = op->getResultTypes()[0].dyn_cast<VectorType>()) {
+ // Map elementwise ops to vec4.
+ SmallVector<int64_t, 4> nativeSize(vecType.getRank() - 1, 1);
+ nativeSize.push_back(4);
+ return nativeSize;
+ }
+ }
return llvm::None;
}
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp
index 864f2d4..7837c17 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp
@@ -304,7 +304,8 @@
patterns.insert<linalg::LinalgTilingPattern<linalg::MatmulOp>,
linalg::LinalgTilingPattern<linalg::FillOp>,
- linalg::LinalgTilingPattern<linalg::BatchMatmulOp>>(
+ linalg::LinalgTilingPattern<linalg::BatchMatmulOp>,
+ linalg::LinalgTilingPattern<linalg::GenericOp>>(
context, tilingOptions,
getLinalgMatchAndReplaceMarker(
{getWorkgroupMemoryMarker(), getWorkgroupMarker()},
@@ -326,7 +327,8 @@
OwningRewritePatternList &patterns) {
patterns.insert<linalg::LinalgVectorizationPattern<linalg::MatmulOp>,
linalg::LinalgVectorizationPattern<linalg::BatchMatmulOp>,
- linalg::LinalgVectorizationPattern<linalg::FillOp>>(
+ linalg::LinalgVectorizationPattern<linalg::FillOp>,
+ linalg::LinalgVectorizationPattern<linalg::GenericOp>>(
context,
linalg::LinalgMarker(Identifier::get(getVectorizeMarker(), context)));
}
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/VectorToGPUPass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/VectorToGPUPass.cpp
index 03cdd3e..bce7f04 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/VectorToGPUPass.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/VectorToGPUPass.cpp
@@ -293,6 +293,44 @@
}
};
+// Lower elementwise operation from N-D vector to 1-D vectors that can be
+// natively supported.
+class ElementwiseLowering : public RewritePattern {
+ public:
+ ElementwiseLowering(MLIRContext *context)
+ : RewritePattern(0, MatchAnyOpTypeTag()) {}
+
+ LogicalResult matchAndRewrite(Operation *op,
+ PatternRewriter &rewriter) const override {
+ if (!op->hasTrait<OpTrait::ElementwiseMappable>() ||
+ op->getNumResults() != 1)
+ return failure();
+ auto vecType = op->getResultTypes()[0].dyn_cast<VectorType>();
+ if (!vecType || vecType.getRank() == 1) return failure();
+
+ SmallVector<Value, 4> newOperands;
+ for (Value operand : op->getOperands()) {
+ if (auto opVecType = operand.getType().dyn_cast<VectorType>()) {
+ auto newType = VectorType::get({opVecType.getNumElements()},
+ opVecType.getElementType());
+ newOperands.push_back(rewriter.create<vector::ShapeCastOp>(
+ op->getLoc(), newType, operand));
+ } else {
+ newOperands.push_back(operand);
+ }
+ }
+ OperationState state(op->getLoc(), op->getName());
+ state.addAttributes(op->getAttrs());
+ state.addOperands(newOperands);
+ state.addTypes({VectorType::get({vecType.getNumElements()},
+ vecType.getElementType())});
+ Operation *newOp = rewriter.createOperation(state);
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, vecType,
+ newOp->getResult(0));
+ return success();
+ }
+};
+
// Lower ExtractStridedSliceOp to an ExtractOp instruction that can be natively
// converted to SPIR-V. Add a BroadcastOp to keep the type consistent, we expect
// the Broadcast to be removed by canonicalization.
@@ -325,7 +363,8 @@
MLIRContext *context) {
OwningRewritePatternList patterns;
patterns.insert<VectorContractLowering, VectorTransferReadToLoad,
- VectorTransferWriteToStore, ExtractStridedLowering>(context);
+ VectorTransferWriteToStore, ExtractStridedLowering,
+ ElementwiseLowering>(context);
applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
}
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/dead_alloc.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/dead_alloc.mlir
new file mode 100644
index 0000000..433b836
--- /dev/null
+++ b/iree/compiler/Conversion/LinalgToSPIRV/test/dead_alloc.mlir
@@ -0,0 +1,20 @@
+// RUN: iree-opt -iree-codegen-optimize-vector-transfer %s | IreeFileCheck %s
+
+module {
+ func @dead_alloc() {
+ %0 = alloc() : memref<8x64xf32, 3>
+ %1 = subview %0[0, 0] [8, 4] [1, 1] : memref<8x64xf32, 3> to
+ memref<8x4xf32, affine_map<(d0, d1) -> (d0 * 64 + d1)>, 3>
+ %c0 = constant 0 : index
+ %cst_0 = constant dense<0.000000e+00> : vector<1x4xf32>
+ vector.transfer_write %cst_0, %1[%c0, %c0] {masked = [false, false]} :
+ vector<1x4xf32>, memref<8x4xf32, affine_map<(d0, d1) -> (d0 * 64 + d1)>, 3>
+ return
+ }
+}
+
+// CHECK-LABEL: func @dead_alloc
+// CHECK-NOT: alloc
+// CHECK-NOT: subview
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/pipeline_test.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/pipeline_test.mlir
index 6c9a4e3..9b2873e 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/test/pipeline_test.mlir
+++ b/iree/compiler/Conversion/LinalgToSPIRV/test/pipeline_test.mlir
@@ -96,3 +96,67 @@
// CHECK-COUNT-12: spv.Load "StorageBuffer" %{{.*}} : vector<4xf32>
// CHECK-COUNT-32: spv.FMul %{{.*}}, %{{.*}} : vector<4xf32>
// CHECK-COUNT-8: spv.Store "StorageBuffer" %{{.*}}, %{{.*}} : vector<4xf32>
+
+// -----
+
+module attributes {
+ spv.target_env =
+ #spv.target_env<#spv.vce<v1.3,
+ [Shader, Float64, Float16, Int64, Int16, Int8, StorageBuffer16BitAccess,
+ StorageUniform16, StoragePushConstant16, StorageBuffer8BitAccess,
+ UniformAndStorageBuffer8BitAccess, StoragePushConstant8, GroupNonUniform,
+ GroupNonUniformVote, GroupNonUniformArithmetic, GroupNonUniformBallot,
+ GroupNonUniformShuffle, GroupNonUniformShuffleRelative, VariablePointers,
+ VariablePointersStorageBuffer],
+ [SPV_KHR_16bit_storage, SPV_KHR_8bit_storage,
+ SPV_KHR_storage_buffer_storage_class, SPV_KHR_variable_pointers]>,
+ ARM:IntegratedGPU,
+ {max_compute_shared_memory_size = 32768 : i32,
+ max_compute_workgroup_invocations = 512 : i32,
+ max_compute_workgroup_size = dense<512> : vector<3xi32>,
+ subgroup_size = 16 : i32}>} {
+ func @matmul_add_fused() attributes {hal.num_workgroups_fn = @matmul_add_fused__num_workgroups__} {
+ %cst = constant 0.000000e+00 : f32
+ %0 = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0, operand_result_index = 3 : i32} : memref<1024x256xf32>
+ %1 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<1024x512xf32>
+ %2 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg1, operand_result_index = 1 : i32} : memref<512x256xf32>
+ %3 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg2, operand_result_index = 2 : i32} : memref<1024x256xf32>
+ %4 = alloc() : memref<1024x256xf32>
+ linalg.fill(%4, %cst) : memref<1024x256xf32>, f32
+ linalg.matmul ins(%1, %2 : memref<1024x512xf32>, memref<512x256xf32>)
+ outs(%4 : memref<1024x256xf32>)
+ linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%4, %3 : memref<1024x256xf32>, memref<1024x256xf32>)
+ outs(%0 : memref<1024x256xf32>) {
+ ^bb0(%arg0: f32, %arg1: f32, %arg2: f32): // no predecessors
+ %5 = addf %arg0, %arg1 : f32
+ linalg.yield %5 : f32
+ }
+ return
+ }
+ func @matmul_add_fused__num_workgroups__
+ (!shapex.ranked_shape<[4096, 4096]>, !shapex.ranked_shape<[4096, 4096]>,
+ !shapex.ranked_shape<[4096, 4096]>) -> (index, index, index)
+ attributes {sym_visibility = "private"}
+ hal.interface @legacy_io attributes {sym_visibility = "private"} {
+ hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+ hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+ hal.interface.binding @arg2, set=0, binding=2, type="StorageBuffer", access="Read"
+ hal.interface.binding @ret0, set=0, binding=3, type="StorageBuffer", access="Write|Discard"
+ }
+}
+
+// CHECK-LABEL: spv.func @matmul_add_fused
+// CHECK-NOT: spv.Store "StorageBuffer"
+// CHECK-NOT: spv.Load "StorageBuffer"
+// CHECK: spv.loop
+// CHECK-COUNT-12: spv.Load "StorageBuffer" %{{.*}} : vector<4xf32>
+// CHECK-COUNT-32: spv.FMul %{{.*}}, %{{.*}} : vector<4xf32>
+// CHECK: spv.mlir.merge
+// CHECK-COUNT-8: spv.Load "StorageBuffer" %{{.*}} : vector<4xf32>
+// CHECK-NOT: spv.Load "StorageBuffer"
+// CHECK-NOT: spv.Store "StorageBuffer"
+// CHECK-COUNT-8: spv.FAdd %{{.*}}, %{{.*}} : vector<4xf32>
+// CHECK-COUNT-8: spv.Store "StorageBuffer" %{{.*}}, %{{.*}} : vector<4xf32>
diff --git a/iree/compiler/Conversion/init_conversions.h b/iree/compiler/Conversion/init_conversions.h
index 05bcdc4..6157e00 100644
--- a/iree/compiler/Conversion/init_conversions.h
+++ b/iree/compiler/Conversion/init_conversions.h
@@ -17,6 +17,7 @@
#include "iree/compiler/Conversion/CodegenUtils/ForOpCanonicalization.h"
#include "iree/compiler/Conversion/HLOToHLO/Passes.h"
+#include "iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensorPasses.h"
#include "iree/compiler/Conversion/HLOToLinalg/Passes.h"
#include "iree/compiler/Conversion/LinalgToLLVM/Passes.h"
#include "iree/compiler/Conversion/LinalgToSPIRV/Passes.h"
diff --git a/iree/compiler/Dialect/Flow/IR/FlowOps.cpp b/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
index cc4e1eb..53b3703 100644
--- a/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
+++ b/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
@@ -17,6 +17,7 @@
#include "iree/compiler/Dialect/Flow/IR/FlowOpUtils.h"
#include "iree/compiler/Dialect/IREE/IR/IREETypes.h"
#include "llvm/ADT/StringExtras.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/Builders.h"
@@ -28,6 +29,7 @@
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/RegionUtils.h"
namespace mlir {
namespace iree_compiler {
@@ -277,27 +279,45 @@
// flow.dispatch.region
//===----------------------------------------------------------------------===//
+/// Inlines operation |op| into the |dispatchRegionOp| by making all operands,
+/// as well as values caputred implicitly by the regions of the operation, that
+/// are outside the dispatch region operands of the dispatch region as well.
+static Operation *inlineOpIntoDispatchRegion(OpBuilder &builder,
+ DispatchRegionOp dispatchRegionOp,
+ Operation *op,
+ BlockAndValueMapping &map) {
+ llvm::SetVector<Value> capturedInputs(op->getOperands().begin(),
+ op->getOperands().end());
+ getUsedValuesDefinedAbove(op->getRegions(), capturedInputs);
+ Block *block = builder.getInsertionBlock();
+ for (Value capturedInput : capturedInputs) {
+ if (map.contains(capturedInput)) continue;
+ dispatchRegionOp.getOperation()->insertOperands(
+ dispatchRegionOp.getOperation()->getNumOperands(), {capturedInput});
+ Value newBlockArgument = block->addArgument(capturedInput.getType());
+ map.map(capturedInput, newBlockArgument);
+ }
+
+ return builder.clone(*op, map);
+}
+
llvm::Optional<std::pair<DispatchRegionOp, Operation *>>
DispatchRegionOp::formFromAnchorOp(Value workload, Operation *anchorOp,
OpBuilder &builder) {
builder.setInsertionPoint(anchorOp);
auto loc = anchorOp->getLoc();
// Map anchor into new dispatch region.
- llvm::SmallVector<Value, 4> capturedInputs(anchorOp->getOperands());
auto drOp = builder.create<DispatchRegionOp>(
loc, llvm::to_vector<1>(anchorOp->getResultTypes()), workload,
- capturedInputs);
+ ArrayRef<Value>());
auto *drBlock = new Block();
drOp.body().push_back(drBlock);
BlockAndValueMapping mapping;
- for (Value capturedInput : capturedInputs) {
- auto blockArg = drBlock->addArgument(capturedInput.getType());
- mapping.map(capturedInput, blockArg);
- }
-
- // Create new body.
builder.setInsertionPointToEnd(drBlock);
- auto *newAnchorOp = builder.clone(*anchorOp, mapping);
+ Operation *newAnchorOp =
+ inlineOpIntoDispatchRegion(builder, drOp, anchorOp, mapping);
+
+ // Insert terminator
builder.create<IREE::Flow::ReturnOp>(loc, newAnchorOp->getResults());
// Replace anchor uses with region result.
@@ -366,16 +386,8 @@
origOpResultValues.push_back(mapping.lookupOrNull(result));
}
- // Add arguments for any op arguments that need to be captured.
- for (Value newArgument : origOp->getOperands()) {
- if (mapping.contains(newArgument)) continue;
- getOperation()->insertOperands(getNumOperands(), {newArgument});
- Value newBlockArgument = block.addArgument(newArgument.getType());
- mapping.map(newArgument, newBlockArgument);
- }
-
- // Clone the op.
- Operation *inlinedOp = builder.clone(*origOp, mapping);
+ Operation *inlinedOp =
+ inlineOpIntoDispatchRegion(builder, *this, origOp, mapping);
// Replace any results from the orig with results from the clone.
for (unsigned i = 0, e = origOp->getNumResults(); i < e; ++i) {
diff --git a/iree/compiler/Dialect/Flow/Transforms/BUILD b/iree/compiler/Dialect/Flow/Transforms/BUILD
index 0a7d576..b1121e9 100644
--- a/iree/compiler/Dialect/Flow/Transforms/BUILD
+++ b/iree/compiler/Dialect/Flow/Transforms/BUILD
@@ -48,6 +48,7 @@
],
deps = [
"//iree/base:signature_mangle",
+ "//iree/compiler/Conversion/HLOToLinalg:HLOToLinalgOnTensors",
"//iree/compiler/Dialect/Flow/Analysis",
"//iree/compiler/Dialect/Flow/Conversion",
"//iree/compiler/Dialect/Flow/Conversion/HLOToFlow",
@@ -61,6 +62,7 @@
"//iree/compiler/Utils",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
+ "@llvm-project//mlir:LinalgOps",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Shape",
"@llvm-project//mlir:ShapeTransforms",
diff --git a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
index 9fd3368..88e17ed 100644
--- a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -44,6 +44,7 @@
DEPS
LLVMSupport
MLIRIR
+ MLIRLinalg
MLIRPass
MLIRShape
MLIRShapeOpsTransforms
@@ -52,6 +53,7 @@
MLIRTransformUtils
MLIRTransforms
iree::base::signature_mangle
+ iree::compiler::Conversion::HLOToLinalg::HLOToLinalgOnTensors
iree::compiler::Dialect::Flow::Analysis
iree::compiler::Dialect::Flow::Conversion
iree::compiler::Dialect::Flow::Conversion::HLOToFlow
diff --git a/iree/compiler/Dialect/Flow/Transforms/DispatchConfig.cpp b/iree/compiler/Dialect/Flow/Transforms/DispatchConfig.cpp
index 1c59969..3581d9b 100644
--- a/iree/compiler/Dialect/Flow/Transforms/DispatchConfig.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/DispatchConfig.cpp
@@ -20,6 +20,7 @@
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#define DEBUG_TYPE "iree-detail"
@@ -129,7 +130,7 @@
}
bool OpDispatchPolicy::isViewModificationOp(Operation *op) {
- return isa<mhlo::ReshapeOp>(op);
+ return isa<mhlo::ReshapeOp, linalg::TensorReshapeOp>(op);
}
int OpDispatchPolicy::getAnchorBenefit(Operation *op) {
@@ -221,8 +222,9 @@
// TODO(b/144530470): replace with tablegen attributes/interfaces.
bool OpDispatchPolicy::isUnsupportedFusionOp(Operation *op) {
- return isa<mhlo::ConcatenateOp, mhlo::ConvOp, mhlo::PadOp, mhlo::ReduceOp,
- mhlo::ReduceWindowOp, mhlo::TorchIndexSelectOp>(op) ||
+ return isa<linalg::IndexedGenericOp, linalg::GenericOp, mhlo::ConcatenateOp,
+ mhlo::ConvOp, mhlo::PadOp, mhlo::ReduceOp, mhlo::ReduceWindowOp,
+ mhlo::TorchIndexSelectOp>(op) ||
(!clEnableConsumerOnlyFusion &&
isa<mhlo::DotOp, mhlo::DotGeneralOp>(op)) ||
isRootOnlyOp(op);
diff --git a/iree/compiler/Dialect/Flow/Transforms/IdentifyDispatchRegions2.cpp b/iree/compiler/Dialect/Flow/Transforms/IdentifyDispatchRegions2.cpp
index f604b39..73e6b6e 100644
--- a/iree/compiler/Dialect/Flow/Transforms/IdentifyDispatchRegions2.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/IdentifyDispatchRegions2.cpp
@@ -20,10 +20,12 @@
#include "iree/compiler/Dialect/Flow/Transforms/DispatchConfig.h"
#include "iree/compiler/Dialect/Flow/Utils/WorkloadUtils.h"
#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SetVector.h"
#include "llvm/Support/Debug.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/RegionUtils.h"
#define DEBUG_TYPE "iree-dispatch"
@@ -271,6 +273,18 @@
}
}
}
+ // Check for values that are used in the region of the op but captured from
+ // outside the region.
+ llvm::SetVector<Value> capturedValues;
+ getUsedValuesDefinedAbove(inlinedOp->getRegions(), capturedValues);
+ for (Value capturedValue : capturedValues) {
+ if (Operation *definingOp = capturedValue.getDefiningOp()) {
+ if (!lastOperandDef ||
+ lastOperandDef.getValue()->isBeforeInBlock(definingOp)) {
+ lastOperandDef = definingOp;
+ }
+ }
+ }
// If the last operand def is already before the dispatch region, there is
// nothing to do.
if (!lastOperandDef ||
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
index d012ba5..6732b18 100644
--- a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -16,13 +16,21 @@
#include <memory>
+#include "iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnTensorPasses.h"
#include "iree/compiler/Dialect/Shape/Conversion/Passes.h"
#include "iree/compiler/Dialect/Shape/Transforms/Passes.h"
#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
#include "mlir/Dialect/Shape/Transforms/Passes.h"
+#include "mlir/Pass/PassOptions.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Transforms/Passes.h"
+static llvm::cl::opt<bool> clEnableLinalgOnTensors(
+ "iree-enable-linalg-on-tensors",
+ llvm::cl::desc(
+ "Enable use of Linalg on tensors for dispatch region creation"),
+ llvm::cl::init(false));
+
namespace mlir {
namespace iree_compiler {
namespace IREE {
@@ -150,6 +158,10 @@
passManager.addNestedPass<FuncOp>(
IREE::Flow::createPrePartitioningConversionPass());
+ if (clEnableLinalgOnTensors) {
+ addHLOToLinalgOnTensorsPasses(passManager);
+ }
+
// First perform module-level analysis that following passes will use to query
// per-function dispatchability information. We run this first so that it only
// needs to run once and will be cached for all of the following passes.
diff --git a/iree/compiler/Dialect/Flow/Transforms/RematerializeDispatchConstants.cpp b/iree/compiler/Dialect/Flow/Transforms/RematerializeDispatchConstants.cpp
index 431f0e3..79cd715 100644
--- a/iree/compiler/Dialect/Flow/Transforms/RematerializeDispatchConstants.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/RematerializeDispatchConstants.cpp
@@ -46,6 +46,8 @@
return true;
} else if (auto value = constantOp.getValue().dyn_cast<DenseElementsAttr>()) {
return value.isSplat();
+ } else if (constantOp.getType().isIntOrFloat()) {
+ return true;
}
// Assume anything unshaped is small. This may not always be true in custom
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/identify_dispatch_regions2_linalg.mlir b/iree/compiler/Dialect/Flow/Transforms/test/identify_dispatch_regions2_linalg.mlir
new file mode 100644
index 0000000..2c82ed7
--- /dev/null
+++ b/iree/compiler/Dialect/Flow/Transforms/test/identify_dispatch_regions2_linalg.mlir
@@ -0,0 +1,50 @@
+// RUN: iree-opt -split-input-file -iree-flow-dispatchability-analysis -iree-flow-identify-dispatch-regions2 %s | IreeFileCheck %s
+
+func @constant_capture(%arg0 : tensor<10x20xf32>) -> tensor<10x20xf32> {
+ %cst1 = constant 1.0 : f32
+ %cst2 = constant dense<2.0> : tensor<10x20xf32>
+ %cst3 = constant dense<
+ [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]> : tensor<10xf32>
+ %0 = linalg.generic
+ {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>,
+ affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0, %cst2, %cst3
+ : tensor<10x20xf32>, tensor<10x20xf32>, tensor<10xf32>) {
+ ^bb0(%arg1 : f32, %arg2 : f32, %arg3 : f32):
+ %1 = addf %arg1, %cst1 : f32
+ %2 = mulf %1, %arg2 : f32
+ %3 = addf %2, %arg3 : f32
+ linalg.yield %3 : f32
+ } -> tensor<10x20xf32>
+ return %0 : tensor<10x20xf32>
+}
+// CHECK: func @constant_capture
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<10x20xf32>
+// CHECK-DAG: %[[CST1:.+]] = constant 1.000000e+00 : f32
+// CHECK-DAG: %[[CST2:.+]] = constant dense<2.000000e+00> : tensor<10x20xf32>
+// CHECK-DAG: %[[CST3:.+]] = constant dense<[1.000000e+00, 2.000000e+00,
+// CHECK-SAME: 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00,
+// CHECK-SAME: 7.000000e+00, 8.000000e+00, 9.000000e+00, 1.000000e+01]>
+// CHECK: %[[RESULT:.+]] = flow.dispatch.region[%{{.+}} : index](
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]] = %[[ARG0]]
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]] = %[[CST2]]
+// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]] = %[[CST3]]
+// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]] = %[[CST1]]
+// CHECK-SAME: ) -> tensor<10x20xf32> {
+// CHECK: %[[RETURN:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[ARG1]], %[[ARG2]], %[[ARG3]]
+// CHECK-SAME: ) {
+// CHECK-NEXT: ^{{[a-zA-Z0-9]+}}(
+// CHECK-SAME: %[[ARG5:.[a-zA-Z0-9_]+]]: f32
+// CHECK-SAME: %[[ARG6:.[a-zA-Z0-9_]+]]: f32
+// CHECK-SAME: %[[ARG7:.[a-zA-Z0-9_]+]]: f32)
+// CHECK: %[[T0:.+]] = addf %[[ARG5]], %[[ARG4]]
+// CHECK: %[[T1:.+]] = mulf %[[T0]], %[[ARG6]]
+// CHECK: %[[T2:.+]] = addf %[[T1]], %[[ARG7]]
+// CHECK: linalg.yield %[[T2]]
+// CHECK: }
+// CHECK: flow.return %[[RETURN]]
+// CHECK: }
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/rematerialize_dispatch_constants.mlir b/iree/compiler/Dialect/Flow/Transforms/test/rematerialize_dispatch_constants.mlir
index 45a0195..0562442 100644
--- a/iree/compiler/Dialect/Flow/Transforms/test/rematerialize_dispatch_constants.mlir
+++ b/iree/compiler/Dialect/Flow/Transforms/test/rematerialize_dispatch_constants.mlir
@@ -65,3 +65,59 @@
}
return %0 : tensor<4x4xf32>
}
+
+// -----
+
+func @constant_capture(%arg0: tensor<10x20xf32>) -> tensor<10x20xf32> {
+ %c200 = constant 200 : index
+ %cst = constant 1.000000e+00 : f32
+ %cst_0 = constant dense<2.000000e+00> : tensor<10x20xf32>
+ %cst_1 = constant dense<
+ [1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00,
+ 6.000000e+00, 7.000000e+00, 8.000000e+00, 9.000000e+00, 1.000000e+01]>
+ : tensor<10xf32>
+ %0 = flow.dispatch.region[%c200 : index]
+ (%arg1 = %arg0 : tensor<10x20xf32>, %arg2 = %cst_0 : tensor<10x20xf32>,
+ %arg3 = %cst_1 : tensor<10xf32>, %arg4 = %cst : f32) -> tensor<10x20xf32> {
+ %1 = linalg.generic
+ {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>,
+ affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg1, %arg2, %arg3
+ : tensor<10x20xf32>, tensor<10x20xf32>, tensor<10xf32>) {
+ ^bb0(%arg5: f32, %arg6: f32, %arg7: f32): // no predecessors
+ %2 = addf %arg5, %arg4 : f32
+ %3 = mulf %2, %arg6 : f32
+ %4 = addf %3, %arg7 : f32
+ linalg.yield %4 : f32
+ } -> tensor<10x20xf32>
+ flow.return %1 : tensor<10x20xf32>
+ }
+ return %0 : tensor<10x20xf32>
+}
+
+// CHECK-LABEL: func @constant_capture
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<10x20xf32>
+// CHECK: %[[CST:.+]] = constant dense<[1.000000e+00, 2.000000e+00,
+// CHECK-SAME: 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00,
+// CHECK-SAME: 7.000000e+00, 8.000000e+00, 9.000000e+00, 1.000000e+01]>
+// CHECK: flow.dispatch.region
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]] = %[[ARG0]]
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]] = %[[CST]]
+// CHECK-DAG: %[[CST_0:.+]] = constant 1.000000e+00 : f32
+// CHECK-DAG: %[[CST_1:.+]] = constant dense<2.000000e+00> : tensor<10x20xf32>
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[ARG1]], %[[CST_1]], %[[ARG2]]
+// CHECK-SAME: ) {
+// CHECK: ^{{[a-zA-Z0-9_]+}}(
+// CHECK-SAME: %[[ARG3:.[a-zA-Z0-9_]+]]: f32
+// CHECK-SAME: %[[ARG4:.[a-zA-Z0-9_]+]]: f32
+// CHECK-SAME: %[[ARG5:.[a-zA-Z0-9_]+]]: f32)
+// CHECK: %[[T0:.+]] = addf %[[ARG3]], %[[CST_0]]
+// CHECK: %[[T1:.+]] = mulf %[[T0]], %[[ARG4]]
+// CHECK: %[[T2:.+]] = addf %[[T1]], %[[ARG5]]
+// CHECK: linalg.yield %[[T2]]
+// CHECK: }
+// CHECK: flow.return %[[RESULT]]
diff --git a/iree/compiler/Dialect/Flow/Utils/DispatchUtils.cpp b/iree/compiler/Dialect/Flow/Utils/DispatchUtils.cpp
index 162c682..7fc8240 100644
--- a/iree/compiler/Dialect/Flow/Utils/DispatchUtils.cpp
+++ b/iree/compiler/Dialect/Flow/Utils/DispatchUtils.cpp
@@ -33,11 +33,11 @@
// TODO(benvanik): replace with op dispatchability interface to allow dialects
// to opt into dispatch.
auto dialectNamespace = op->getDialect()->getNamespace();
- return dialectNamespace == mhlo::MhloDialect::getDialectNamespace() ||
+ return dialectNamespace == FlowDialect::getDialectNamespace() ||
+ dialectNamespace == linalg::LinalgDialect::getDialectNamespace() ||
+ dialectNamespace == mhlo::MhloDialect::getDialectNamespace() ||
dialectNamespace == mlir::StandardOpsDialect::getDialectNamespace() ||
- dialectNamespace == FlowDialect::getDialectNamespace() ||
- dialectNamespace == ShapeDialect::getDialectNamespace() ||
- dialectNamespace == linalg::LinalgDialect::getDialectNamespace();
+ dialectNamespace == ShapeDialect::getDialectNamespace();
}
namespace {
diff --git a/iree/compiler/Dialect/HAL/IR/HALOps.cpp b/iree/compiler/Dialect/HAL/IR/HALOps.cpp
index 463b49e..9bf155b 100644
--- a/iree/compiler/Dialect/HAL/IR/HALOps.cpp
+++ b/iree/compiler/Dialect/HAL/IR/HALOps.cpp
@@ -1356,8 +1356,9 @@
}
OptionalParseResult parseResult = parser.parseOptionalRegion(*body);
- if (parseResult.hasValue() && failed(*parseResult))
+ if (parseResult.hasValue() && failed(*parseResult)) {
return failure();
+ }
// Ensure that this module has a valid terminator.
ExecutableTargetOp::ensureTerminator(*body, parser.getBuilder(),
@@ -1410,8 +1411,9 @@
return failure();
}
OptionalParseResult parseResult = parser.parseOptionalRegion(*body);
- if (parseResult.hasValue() && failed(*parseResult))
+ if (parseResult.hasValue() && failed(*parseResult)) {
return failure();
+ }
// Ensure that this module has a valid terminator.
ExecutableBinaryOp::ensureTerminator(*body, parser.getBuilder(),
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/ConstantEncoder.cpp b/iree/compiler/Dialect/VM/Target/Bytecode/ConstantEncoder.cpp
index 64e1079..faf41fe 100644
--- a/iree/compiler/Dialect/VM/Target/Bytecode/ConstantEncoder.cpp
+++ b/iree/compiler/Dialect/VM/Target/Bytecode/ConstantEncoder.cpp
@@ -106,6 +106,19 @@
return flatbuffers_uint8_vec_end(fbb);
}
+static flatbuffers_uint8_vec_ref_t serializeConstantF16Array(
+ DenseFPElementsAttr attr, FlatbufferBuilder &fbb) {
+ flatbuffers_uint8_vec_start(fbb);
+ uint8_t *bytePtr = flatbuffers_uint8_vec_extend(
+ fbb, attr.getNumElements() * sizeof(uint16_t));
+ uint16_t *nativePtr = reinterpret_cast<uint16_t *>(bytePtr);
+ for (const APFloat &value : attr.getFloatValues()) {
+ *(nativePtr++) =
+ value.bitcastToAPInt().extractBitsAsZExtValue(16, 0) & UINT16_MAX;
+ }
+ return flatbuffers_uint8_vec_end(fbb);
+}
+
flatbuffers_uint8_vec_ref_t serializeConstant(Location loc,
ElementsAttr elementsAttr,
FlatbufferBuilder &fbb) {
@@ -126,6 +139,8 @@
}
} else if (auto attr = elementsAttr.dyn_cast<DenseFPElementsAttr>()) {
switch (attr.getType().getElementTypeBitWidth()) {
+ case 16:
+ return serializeConstantF16Array(attr, fbb);
case 32:
return serializeConstantF32Array(attr, fbb);
case 64:
diff --git a/iree/compiler/Dialect/VM/Target/Bytecode/test/constant_encoding.mlir b/iree/compiler/Dialect/VM/Target/Bytecode/test/constant_encoding.mlir
index fc4d1a2..78298bb 100644
--- a/iree/compiler/Dialect/VM/Target/Bytecode/test/constant_encoding.mlir
+++ b/iree/compiler/Dialect/VM/Target/Bytecode/test/constant_encoding.mlir
@@ -47,4 +47,15 @@
// CHECK-NEXT: 63
// CHECK-NEXT: ]
vm.rodata @splat_float32s dense<1.000000e+00> : tensor<3xf32>
+
+ // CHECK: "data": [
+ // CHECK-NEXT: 0,
+ // CHECK-NEXT: 60,
+ // CHECK-NEXT: 0,
+ // CHECK-NEXT: 64,
+ // CHECK-NEXT: 0,
+ // CHECK-NEXT: 66
+ // CHECK-NEXT: ]
+ vm.rodata @dense_float16s dense<[1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<3xf16>
+
}
diff --git a/iree/test/e2e/models/edge_detection.mlir b/iree/test/e2e/models/edge_detection.mlir
index bb8b47a..c290e84 100644
--- a/iree/test/e2e/models/edge_detection.mlir
+++ b/iree/test/e2e/models/edge_detection.mlir
@@ -1,6 +1,8 @@
// RUN: iree-run-mlir -export-all -iree-hal-target-backends=vmla %s -function-input="1x128x128x1xf32" | IreeFileCheck %s
// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=llvm-ir %s -function-input="1x128x128x1xf32" | IreeFileCheck %s)
// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=vulkan-spirv %s -function-input="1x128x128x1xf32" | IreeFileCheck %s)
+// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=llvm-ir -iree-enable-linalg-on-tensors %s -function-input="1x128x128x1xf32" | IreeFileCheck %s)
+// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=vulkan-spirv -iree-enable-linalg-on-tensors %s -function-input="1x128x128x1xf32" | IreeFileCheck %s)
// Image edge detection module generated by iree/colab/edge_detection.ipynb.
//
diff --git a/iree/test/e2e/models/fragment_000.mlir b/iree/test/e2e/models/fragment_000.mlir
index dd9f902..8328f86 100644
--- a/iree/test/e2e/models/fragment_000.mlir
+++ b/iree/test/e2e/models/fragment_000.mlir
@@ -1,6 +1,8 @@
// RUN: iree-run-mlir -export-all -iree-hal-target-backends=vmla %s -function-input="f32=0" -function-input="5x1xf32=[1][-2][-3][4][-5]" -function-input="f32=1" -function-input="5x5xf32=[3.46499 -7.64389 -5.72249 5.98053 17.6892][2.9707 -6.20734 -4.25962 4.76055 13.8784][2.47641 -4.77079 -2.79675 3.54056 10.0675][1.98212 -3.33424 -1.33388 2.32058 6.25666][1.48783 -1.8977 0.12899 1.1006 2.4458]" -function-input="5xf32=0 0 0 0 0" | IreeFileCheck %s --implicit-check-not="[" --implicit-check-not="]"
// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=llvm-ir %s -function-input="f32=0" -function-input="5x1xf32=[1][-2][-3][4][-5]" -function-input="f32=1" -function-input="5x5xf32=[3.46499 -7.64389 -5.72249 5.98053 17.6892][2.9707 -6.20734 -4.25962 4.76055 13.8784][2.47641 -4.77079 -2.79675 3.54056 10.0675][1.98212 -3.33424 -1.33388 2.32058 6.25666][1.48783 -1.8977 0.12899 1.1006 2.4458]" -function-input="5xf32=0 0 0 0 0" | IreeFileCheck %s --implicit-check-not="[" --implicit-check-not="]")
// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=vulkan-spirv %s -function-input="f32=0" -function-input="5x1xf32=[1][-2][-3][4][-5]" -function-input="f32=1" -function-input="5x5xf32=[3.46499 -7.64389 -5.72249 5.98053 17.6892][2.9707 -6.20734 -4.25962 4.76055 13.8784][2.47641 -4.77079 -2.79675 3.54056 10.0675][1.98212 -3.33424 -1.33388 2.32058 6.25666][1.48783 -1.8977 0.12899 1.1006 2.4458]" -function-input="5xf32=0 0 0 0 0" | IreeFileCheck %s --implicit-check-not="[" --implicit-check-not="]")
+// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=llvm-ir -iree-enable-linalg-on-tensors %s -function-input="f32=0" -function-input="5x1xf32=[1][-2][-3][4][-5]" -function-input="f32=1" -function-input="5x5xf32=[3.46499 -7.64389 -5.72249 5.98053 17.6892][2.9707 -6.20734 -4.25962 4.76055 13.8784][2.47641 -4.77079 -2.79675 3.54056 10.0675][1.98212 -3.33424 -1.33388 2.32058 6.25666][1.48783 -1.8977 0.12899 1.1006 2.4458]" -function-input="5xf32=0 0 0 0 0" | IreeFileCheck %s --implicit-check-not="[" --implicit-check-not="]")
+// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=vulkan-spirv -iree-enable-linalg-on-tensors %s -function-input="f32=0" -function-input="5x1xf32=[1][-2][-3][4][-5]" -function-input="f32=1" -function-input="5x5xf32=[3.46499 -7.64389 -5.72249 5.98053 17.6892][2.9707 -6.20734 -4.25962 4.76055 13.8784][2.47641 -4.77079 -2.79675 3.54056 10.0675][1.98212 -3.33424 -1.33388 2.32058 6.25666][1.48783 -1.8977 0.12899 1.1006 2.4458]" -function-input="5xf32=0 0 0 0 0" | IreeFileCheck %s --implicit-check-not="[" --implicit-check-not="]")
// CHECK-LABEL: EXEC @main_entry_dispatch_3
func @main_entry_dispatch_3(
diff --git a/iree/test/e2e/models/fullyconnected.mlir b/iree/test/e2e/models/fullyconnected.mlir
index b2ba5f4..ab9c466 100644
--- a/iree/test/e2e/models/fullyconnected.mlir
+++ b/iree/test/e2e/models/fullyconnected.mlir
@@ -1,6 +1,8 @@
// RUN: iree-run-mlir -export-all %s -iree-hal-target-backends=vmla -function-input="1x5xf32=1,-2,-3,4,-5" -function-input="1x5x3x1xf32=15,14,13,12,11,10,9,8,7,6,5,4,3,2,1" | IreeFileCheck %s
// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir -export-all %s -iree-hal-target-backends=llvm-ir -function-input="1x5xf32=1,-2,-3,4,-5" -function-input="1x5x3x1xf32=15,14,13,12,11,10,9,8,7,6,5,4,3,2,1" -iree-enable-consumer-only-fusion | IreeFileCheck %s)
// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir -export-all %s -iree-hal-target-backends=vulkan-spirv -function-input="1x5xf32=1,-2,-3,4,-5" -function-input="1x5x3x1xf32=15,14,13,12,11,10,9,8,7,6,5,4,3,2,1" | IreeFileCheck %s)
+// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir -export-all %s -iree-hal-target-backends=llvm-ir -function-input="1x5xf32=1,-2,-3,4,-5" -function-input="1x5x3x1xf32=15,14,13,12,11,10,9,8,7,6,5,4,3,2,1" -iree-enable-linalg-on-tensors | IreeFileCheck %s)
+// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir -export-all %s -iree-hal-target-backends=vulkan-spirv -function-input="1x5xf32=1,-2,-3,4,-5" -function-input="1x5x3x1xf32=15,14,13,12,11,10,9,8,7,6,5,4,3,2,1" -iree-enable-linalg-on-tensors | IreeFileCheck %s)
// CHECK-LABEL: EXEC @main
func @main(%arg0: tensor<1x5xf32>, %arg1: tensor<1x5x3x1xf32>) -> tuple<tensor<5x1x5xf32>>
diff --git a/iree/test/e2e/models/mnist_fake_weights.mlir b/iree/test/e2e/models/mnist_fake_weights.mlir
index 8404532..380899d 100644
--- a/iree/test/e2e/models/mnist_fake_weights.mlir
+++ b/iree/test/e2e/models/mnist_fake_weights.mlir
@@ -3,6 +3,8 @@
// RUN: iree-run-mlir -export-all -iree-hal-target-backends=vmla %s -function-input="1x28x28x1xf32" | IreeFileCheck %s
// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=llvm-ir %s -function-input="1x28x28x1xf32" | IreeFileCheck %s)
// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=vulkan-spirv %s -function-input="1x28x28x1xf32" | IreeFileCheck %s)
+// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=llvm-ir -iree-enable-linalg-on-tensors %s -function-input="1x28x28x1xf32" | IreeFileCheck %s)
+// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir -export-all -iree-hal-target-backends=vulkan-spirv -iree-enable-linalg-on-tensors %s -function-input="1x28x28x1xf32" | IreeFileCheck %s)
module {
flow.variable @"__iree_flow___sm_node17__model.layer-1.kernel" dense<1.000000e+00> : tensor<784x128xf32> attributes {noinline, sym_visibility = "private"}
diff --git a/iree/test/e2e/models/unidirectional_lstm.mlir b/iree/test/e2e/models/unidirectional_lstm.mlir
index 510eb4a..302f79e 100644
--- a/iree/test/e2e/models/unidirectional_lstm.mlir
+++ b/iree/test/e2e/models/unidirectional_lstm.mlir
@@ -3,6 +3,8 @@
// RUN: iree-run-mlir %s -iree-hal-target-backends=vmla -function-input="1x5xf32=[0 1 0 3 4]" -function-input="1x5x2x2xf32=[1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20]" | IreeFileCheck %s --implicit-check-not="[" --implicit-check-not="]"
// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir %s -iree-hal-target-backends=llvm-ir -function-input="1x5xf32=[0 1 0 3 4]" -function-input="1x5x2x2xf32=[1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20]" | IreeFileCheck %s --implicit-check-not="[" --implicit-check-not="]")
// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir %s -iree-hal-target-backends=vulkan-spirv -function-input="1x5xf32=[0 1 0 3 4]" -function-input="1x5x2x2xf32=[1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20]" | IreeFileCheck %s --implicit-check-not="[" --implicit-check-not="]")
+// RUN: [[ $IREE_LLVMJIT_DISABLE == 1 ]] || (iree-run-mlir %s -iree-hal-target-backends=llvm-ir -function-input="1x5xf32=[0 1 0 3 4]" -function-input="1x5x2x2xf32=[1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20]" -iree-enable-linalg-on-tensors | IreeFileCheck %s --implicit-check-not="[" --implicit-check-not="]")
+// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir %s -iree-hal-target-backends=vulkan-spirv -function-input="1x5xf32=[0 1 0 3 4]" -function-input="1x5x2x2xf32=[1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20]" -iree-enable-linalg-on-tensors | IreeFileCheck %s --implicit-check-not="[" --implicit-check-not="]")
// Exported via the XLA HLO Importer
// The resulting MLIR was modified by hand by changing all large constants to be
diff --git a/iree/tools/run_lit.sh b/iree/tools/run_lit.sh
index 704c7cc..78cbf66 100755
--- a/iree/tools/run_lit.sh
+++ b/iree/tools/run_lit.sh
@@ -17,9 +17,16 @@
set -o pipefail
if [ -z "${RUNFILES_DIR}" ]; then
- # Some versions of bazel do not set RUNFILES_DIR. Instead they just cd
- # into the directory.
- RUNFILES_DIR="$PWD"
+ if [ -f "CMakeCache.txt" ]; then
+ # If running under CMake/CTest in the build directory, just scope to the
+ # iree directory to avoid blowing up the search through things like
+ # bazel out directories and the like.
+ RUNFILES_DIR="$PWD/iree"
+ else
+ # Some versions of bazel do not set RUNFILES_DIR. Instead they just cd
+ # into the directory.
+ RUNFILES_DIR="$PWD"
+ fi
fi
# Detect whether cygwin/msys2 paths need to be translated.
diff --git a/packaging/python/.gitignore b/packaging/python/.gitignore
deleted file mode 100644
index b357fe2..0000000
--- a/packaging/python/.gitignore
+++ /dev/null
@@ -1,4 +0,0 @@
-build/
-dist/
-*.egg-info
-
diff --git a/packaging/python/BUILD.bazel b/packaging/python/BUILD.bazel
deleted file mode 100644
index 4e48edd..0000000
--- a/packaging/python/BUILD.bazel
+++ /dev/null
@@ -1,39 +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.
-
-package(
- features = ["layering_check"],
- licenses = ["notice"],
-)
-
-# This is a dummy binary that has the side-effect of building all of the TF
-# python bindings. It is used to build wheel files.
-py_binary(
- name = "all_pyiree_packages",
- srcs = ["dummy_exclude_from_package.py"],
- legacy_create_init = False,
- main = "dummy_exclude_from_package.py",
- python_version = "PY3",
- tags = [
- # Do not build with ... expansion
- "manual",
- ],
- deps = [
- "//bindings/python:pathsetup", # build_cleaner: keep
- "//bindings/python/pyiree/compiler", # build_cleaner: keep
- "//bindings/python/pyiree/rt", # build_cleaner: keep
- "//integrations/tensorflow/bindings/python/pyiree/tf/compiler", # build_cleaner: keep
- "//integrations/tensorflow/bindings/python/pyiree/tf/support", # build_cleaner: keep
- ],
-)
diff --git a/packaging/python/README.md b/packaging/python/README.md
deleted file mode 100644
index 3d387a2..0000000
--- a/packaging/python/README.md
+++ /dev/null
@@ -1,87 +0,0 @@
-# Python packaging scripts.
-
-Note that packages will be placed in `packaging/python/dist` with the canonical
-instructions. However, the setup scripts can be run from anywhere and will
-create `build` and `dist` directories where run. Wheels can be installed with
-`pip3 install --user dist/*.whl`.
-
-## Building core wheels with CMake
-
-Most of IREE is built/packaged with CMake. For the parts that build with CMake,
-this is preferred.
-
-Canonical instructions follow:
-
-### Linux
-
-```shell
-export LDFLAGS=-fuse-ld=/usr/bin/ld.lld
-export PYIREE_CMAKE_BUILD_ROOT="${HOME?}/build-iree-release"
-export IREE_SRC="${HOME?}/src/iree"
-rm -Rf "${PYIREE_CMAKE_BUILD_ROOT?}"; mkdir -p "${PYIREE_CMAKE_BUILD_ROOT?}"
-cmake -GNinja -B"${PYIREE_CMAKE_BUILD_ROOT?}" -H"${IREE_SRC}" \
- -DCMAKE_BUILD_TYPE=Release \
- -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ \
- -DIREE_BUILD_PYTHON_BINDINGS=ON -DIREE_BUILD_SAMPLES=OFF
-(cd "${PYIREE_CMAKE_BUILD_ROOT?}" && ninja)
-(cd "${IREE_SRC?}/packaging/python" && (
- rm -Rf build;
- python3 setup_compiler.py bdist_wheel;
- rm -Rf build;
- python3 setup_rt.py bdist_wheel))
-```
-
-## Building IREE/TensorFlow wheels
-
-If building TensorFlow integration wheels, then this must be done via Bazel. In
-this case, it can be easiest to just package everything from a Bazel build to
-avoid multiple steps.
-
-Canonical instructions follow:
-
-### Env Setup
-
-```shell
-IREE_SRC=$HOME/src/iree
-export PYIREE_BAZEL_BUILD_ROOT="$IREE_SRC/bazel-bin"
-if which cygpath; then
- export PYIREE_BAZEL_BUILD_ROOT="$(cygpath -w "$PYIREE_BAZEL_BUILD_ROOT")"
-fi
-```
-
-### Building:
-
-Optionally add: `--define=PYIREE_TF_DISABLE_KERNELS=1` to build a 'thin' (less
-functional) version without TensorFlow kernels. This should not be done for
-released binaries but can help while developing.
-
-Note that bazel does not always build properly named artifacts. See the tool
-`hack_python_package_from_runfiles.py` to extract and fixup artifacts from a
-bazel-bin directory. If using this mechanism, then the environment variable
-`PYIREE_PYTHON_ROOT` should be set to a suitable temp directory.
-
-```shell
-cd $IREE_SRC
-bazel build -c opt \
- //packaging/python:all_pyiree_packages
-```
-
-# Packaging
-
-```shell
-(cd $IREE_SRC/packaging/python && (
- rm -Rf build;
- python3 setup_tf.py bdist_wheel))
-```
-
-```shell
-(cd $IREE_SRC/packaging/python && (
- rm -Rf build;
- python3 setup_compiler.py bdist_wheel))
-```
-
-```shell
-(cd $IREE_SRC/packaging/python && (
- rm -Rf build;
- python3 setup_rt.py bdist_wheel))
-```
diff --git a/packaging/python/__init__.py b/packaging/python/__init__.py
deleted file mode 100644
index 8b13789..0000000
--- a/packaging/python/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-
diff --git a/packaging/python/common_setup.py b/packaging/python/common_setup.py
deleted file mode 100644
index 727ce24..0000000
--- a/packaging/python/common_setup.py
+++ /dev/null
@@ -1,132 +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.
-
-import os
-import platform
-import setuptools
-import sys
-import sysconfig
-from datetime import date
-
-
-def get_exe_suffix():
- if platform.system() == "Windows":
- return ".exe"
- else:
- return ""
-
-
-def get_package_dir(prefix=("bindings", "python")):
- explicit_root = os.environ.get("PYIREE_PYTHON_ROOT")
- if explicit_root:
- return explicit_root
-
- # Use env variables based on build system type.
- cmake_build_root = os.environ.get("PYIREE_CMAKE_BUILD_ROOT")
- bazel_build_root = os.environ.get("PYIREE_BAZEL_BUILD_ROOT")
-
- if cmake_build_root and bazel_build_root:
- print("ERROR: Both PYIREE_CMAKE_BUILD_ROOT and PYIREE_BAZEL_BUILD_ROOT"
- " cannot be set at the same time")
- sys.exit(1)
-
- if cmake_build_root:
- print("Using CMake build root:", cmake_build_root)
- pkg_dir = os.path.join(cmake_build_root, *prefix)
- elif bazel_build_root:
- print("Using Bazel build root:", bazel_build_root)
- if not os.path.isdir(bazel_build_root):
- print("ERROR: Could not find bazel-bin:", bazel_build_root)
- sys.exit(1)
- # Find the path to the runfiles of the built target:
- # //bindings/python/packaging:all_pyiree_packages
- runfiles_dir = os.path.join(
- bazel_build_root, "packaging", "python",
- "all_pyiree_packages%s.runfiles" % (get_exe_suffix(),))
- if not os.path.isdir(runfiles_dir):
- print("ERROR: Could not find build target 'all_pyiree_packages':",
- runfiles_dir)
- print("Make sure to build target",
- "//packaging/python:all_pyiree_packages")
- sys.exit(1)
- # And finally seek into the corresponding path in the runfiles dir.
- # Aren't bazel paths fun???
- # Note that the "iree_core" path segment corresponds to the workspace name.
- pkg_dir = os.path.join(runfiles_dir, "iree_core", *prefix)
- else:
- print("ERROR: No build directory specified. Set one of these variables:")
- print(" PYIREE_CMAKE_BUILD_ROOT=/path/to/cmake/build")
- sys.exit(1)
-
- if not os.path.exists(pkg_dir):
- print("ERROR: Package path does not exist:", pkg_dir)
- sys.exit(1)
- return pkg_dir
-
-
-def get_default_date_version():
- today = date.today()
- return today.strftime("%Y%m%d")
-
-
-def get_setup_defaults(sub_project, description, package_dir=None):
- if not package_dir:
- package_dir = get_package_dir()
- return {
- "name": "google-iree-%s" % (sub_project,),
- "version": get_default_date_version(),
- "author": "The IREE Team at Google",
- "author_email": "iree-discuss@googlegroups.com",
- "description": description,
- "long_description": description,
- "long_description_content_type": "text/plain",
- "url": "https://github.com/google/iree",
- "package_dir": {
- "": package_dir,
- },
- "classifiers": [
- "Programming Language :: Python :: 3",
- "License :: OSI Approved :: Apache License",
- "Operating System :: OS Independent",
- "Development Status :: 3 - Alpha",
- ],
- "python_requires": ">=3.6",
- }
-
-
-def setup(**kwargs):
- # See: https://stackoverflow.com/q/45150304
- try:
- from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
-
- class bdist_wheel(_bdist_wheel):
-
- def finalize_options(self):
- _bdist_wheel.finalize_options(self)
- self.root_is_pure = False
- except ImportError:
- bdist_wheel = None
-
- # Need to include platform specific extensions binaries:
- # Windows: .pyd
- # macOS: .dylib
- # Other: .so
- # Unfortunately, bazel is imprecise and scatters .so files around, so
- # need to be specific.
- package_data = {
- "": ["*%s" % (sysconfig.get_config_var("EXT_SUFFIX"),)],
- }
- setuptools.setup(package_data=package_data,
- cmdclass={"bdist_wheel": bdist_wheel},
- **kwargs)
diff --git a/packaging/python/hack_python_package_from_runfiles.py b/packaging/python/hack_python_package_from_runfiles.py
deleted file mode 100644
index 61c2d6e..0000000
--- a/packaging/python/hack_python_package_from_runfiles.py
+++ /dev/null
@@ -1,102 +0,0 @@
-#!/usr/bin/python
-# 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.
-
-# Given a runfiles directory from a bazel build, does surgery to extract
-# a usable python package directory. In addition to the bazel directory
-# structure being unnecessarily obtuse, it is also really hard to actually
-# name files correctly. This affects python extension modules which must be
-# named with a specific extension suffix. Bazel is extremely unflexible and
-# we patch around it with this script. For the record, there are various ways
-# to write custom rules to do this more natively, but it is all complicated
-# and needless complexity. We opt for a script that is at least readable by
-# mere mortals and in one place.
-# Usage:
-# ./this_script <dest_dir> <path to bazel-bin>
-
-import os
-import platform
-import shutil
-import sys
-import sysconfig
-
-FILE_NAME_MAP = {
- "binding.so": "binding{}".format(sysconfig.get_config_var("EXT_SUFFIX")),
- "binding.pyd": False,
- "binding.dylib": False,
-}
-
-
-def get_exe_suffix():
- if platform.system() == "Windows":
- return ".exe"
- else:
- return ""
-
-
-def copy_prefix(dest_dir, runfiles_dir, prefix):
- # And finally seek into the corresponding path in the runfiles dir.
- # Aren't bazel paths fun???
- # Note that the "iree_core" path segment corresponds to the workspace name.
- pkg_dir = os.path.join(runfiles_dir, "iree_core", *prefix)
- if not os.path.exists(pkg_dir):
- return
- dest_dir = os.path.join(dest_dir)
- for root, dirs, files in os.walk(pkg_dir):
- assert root.startswith(pkg_dir)
- dest_prefix = root[len(pkg_dir):]
- if dest_prefix.startswith(os.path.sep):
- dest_prefix = dest_prefix[1:]
- local_dest_dir = os.path.join(dest_dir, dest_prefix)
- os.makedirs(local_dest_dir, exist_ok=True)
- for file in files:
- copy_file(os.path.join(root, file), local_dest_dir)
-
-
-def copy_file(src_file, dst_dir):
- basename = os.path.basename(src_file)
- dst_file = os.path.join(dst_dir, basename)
- mapped_name = FILE_NAME_MAP.get(basename)
- if mapped_name is False:
- # Skip.
- return
- elif mapped_name is not None:
- dst_file = os.path.join(dst_dir, mapped_name)
- shutil.copyfile(src_file, dst_file, follow_symlinks=True)
-
-
-def main():
- # Parse args.
- dest_dir = sys.argv[1]
- bazel_bin = sys.argv[2] if len(sys.argv) > 2 else os.path.join(
- os.path.dirname(__file__), "..", "..", "bazel-bin")
-
- # Find the path to the runfiles of the built target:
- # //bindings/python/packaging:all_pyiree_packages
- runfiles_dir = os.path.join(
- bazel_bin, "packaging", "python",
- "all_pyiree_packages%s.runfiles" % (get_exe_suffix(),))
- if not os.path.isdir(runfiles_dir):
- print("ERROR: Could not find build target 'all_pyiree_packages':",
- runfiles_dir)
- print("Make sure to build target", "//packaging/python:all_pyiree_packages")
- sys.exit(1)
-
- copy_prefix(dest_dir, runfiles_dir, ("bindings", "python"))
- copy_prefix(dest_dir, runfiles_dir,
- ("integrations", "tensorflow", "bindings", "python"))
-
-
-if __name__ == "__main__":
- main()
diff --git a/packaging/python/setup_compiler.py b/packaging/python/setup_compiler.py
deleted file mode 100644
index 57f4479..0000000
--- a/packaging/python/setup_compiler.py
+++ /dev/null
@@ -1,51 +0,0 @@
-#!/usr/bin/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.
-
-# Build platform specific wheel files for the pyiree.rt package.
-# Built artifacts are per-platform and build out of the build tree.
-# Usage:
-# ------
-# Windows with CMake:
-# export CMAKE_BUILD_ROOT='D:\src\build-iree' # Must be native path
-# python ./setup_compiler.py bdist_wheel
-
-import os
-import setuptools
-import sys
-
-# Ensure that path starts here for execution as a script.
-sys.path.insert(0, os.path.dirname(__file__))
-import common_setup
-
-
-def run():
- packages = setuptools.find_namespace_packages(
- common_setup.get_package_dir(),
- include=["pyiree.compiler", "pyiree.compiler.*"],
- exclude=["*.CMakeFiles"])
- print("Found packages:", packages)
- setup_kwargs = common_setup.get_setup_defaults(
- sub_project="compiler", description="IREE Generic Compiler")
- common_setup.setup(packages=packages,
- ext_modules=[
- setuptools.Extension(name="pyiree.compiler.binding",
- sources=[]),
- ],
- **setup_kwargs)
-
-
-if __name__ == "__main__":
- run()
diff --git a/packaging/python/setup_rt.py b/packaging/python/setup_rt.py
deleted file mode 100644
index c9fc948..0000000
--- a/packaging/python/setup_rt.py
+++ /dev/null
@@ -1,47 +0,0 @@
-#!/usr/bin/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.
-
-# Build platform specific wheel files for the pyiree.rt package.
-# Built artifacts are per-platform and build out of the build tree.
-
-import os
-import setuptools
-import sys
-
-# Ensure that path starts here for execution as a script.
-sys.path.insert(0, os.path.dirname(__file__))
-import common_setup
-
-
-def run():
- packages = setuptools.find_namespace_packages(
- common_setup.get_package_dir(),
- include=["pyiree.rt", "pyiree.rt.*"],
- exclude=["*.CMakeFiles"])
- print("Found packages:", packages)
- setup_kwargs = common_setup.get_setup_defaults(
- sub_project="rt",
- description="IREE Runtime Components (for executing compiled programs)")
- common_setup.setup(packages=packages,
- ext_modules=[
- setuptools.Extension(name="pyiree.rt.binding",
- sources=[]),
- ],
- **setup_kwargs)
-
-
-if __name__ == "__main__":
- run()
diff --git a/packaging/python/setup_tf.py b/packaging/python/setup_tf.py
deleted file mode 100644
index 252c756..0000000
--- a/packaging/python/setup_tf.py
+++ /dev/null
@@ -1,59 +0,0 @@
-#!/usr/bin/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.
-
-# Build platform specific wheel files for the pyiree.tf packages.
-# Built artifacts are per-platform and build out of the build tree.
-
-import os
-import platform
-import setuptools
-import sys
-
-# Ensure that path starts here for execution as a script.
-sys.path.insert(0, os.path.dirname(__file__))
-import common_setup
-
-
-def run():
- package_dir = common_setup.get_package_dir(prefix=("integrations",
- "tensorflow", "bindings",
- "python"))
- packages = setuptools.find_namespace_packages(package_dir,
- include=[
- "pyiree.tf.compiler",
- "pyiree.tf.compiler.*",
- "pyiree.tf.support",
- "pyiree.tf.support.*"
- ],
- exclude=["*.CMakeFiles"])
- print("Found packages:", packages)
- if not packages:
- print("ERROR: Did not find packages under", package_dir)
- sys.exit(1)
- setup_kwargs = common_setup.get_setup_defaults(
- sub_project="tf",
- description="IREE TensorFlow Compiler",
- package_dir=package_dir)
- common_setup.setup(packages=packages,
- ext_modules=[
- setuptools.Extension(name="pyiree.tf.compiler.binding",
- sources=[]),
- ],
- **setup_kwargs)
-
-
-if __name__ == "__main__":
- run()
