Merge pull request #3594 from rsuderman:main-to-google
PiperOrigin-RevId: 338738884
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index d5d5bb0..7e91cc7 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -50,8 +50,8 @@
uses: actions/checkout@v2
- name: Setting up python
uses: actions/setup-python@v2
- # We have to explicitly fetch the base branch as well
- name: Fetching Base Branch
+ # We have to explicitly fetch the base branch as well
run: git fetch --no-tags --prune --depth=1 origin "${BASE_REF?}:${BASE_REF?}"
- name: Install yapf
run: python3 -m pip install yapf
@@ -65,6 +65,26 @@
printf "You can fix the lint errors above by running\n"
printf " git diff -U0 "${BASE_REF?}" | python3 third_party/format_diff/format_diff.py yapf -i\n"
+ pytype:
+ runs-on: ubuntu-18.04
+ env:
+ BASE_REF: ${{ github.base_ref }}
+ steps:
+ - name: Checking out repository
+ uses: actions/checkout@v2
+ - name: Setting up python
+ uses: actions/setup-python@v2
+ with:
+ # Pytype does not support python3.9, which this action defaults to.
+ python-version: '3.8'
+ - name: Fetching Base Branch
+ # We have to explicitly fetch the base branch as well
+ run: git fetch --no-tags --prune --depth=1 origin "${BASE_REF?}:${BASE_REF?}"
+ - name: Install pytype
+ run: python3 -m pip install pytype
+ - name: Run pytype on changed files
+ run: ./build_tools/pytype/check_diff.sh "${BASE_REF?}"
+
clang-format:
runs-on: ubuntu-18.04
env:
@@ -79,8 +99,8 @@
chmod +x /tmp/git-clang-format
- name: Checking out repository
uses: actions/checkout@v2
- # We have to explicitly fetch the base branch as well
- name: Fetching Base Branch
+ # We have to explicitly fetch the base branch as well
run: git fetch --no-tags --prune --depth=1 origin "${BASE_REF?}:${BASE_REF?}"
- name: Running clang-format on changed source files
run: |
@@ -102,8 +122,8 @@
steps:
- name: Checking out repository
uses: actions/checkout@v2
- # We have to explicitly fetch the base branch as well
- name: Fetching Base Branch
+ # We have to explicitly fetch the base branch as well
run: git fetch --no-tags --prune --depth=1 origin "${BASE_REF?}:${BASE_REF?}"
- name: Checking tabs
run: ./scripts/check_tabs.sh "${BASE_REF?}"
diff --git a/.gitignore b/.gitignore
index 96a5b66..b0a8cbc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
# Python
*.pyc
**/.ipynb_checkpoints/
+.pytype/
# Visual Studio files
.vs/
diff --git a/build_tools/pytype/check_diff.sh b/build_tools/pytype/check_diff.sh
new file mode 100755
index 0000000..2b9851a
--- /dev/null
+++ b/build_tools/pytype/check_diff.sh
@@ -0,0 +1,96 @@
+#!/bin/bash
+# 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.
+
+# Uses git diff to run pytype on changed files.
+# Example Usage:
+# Defaults to comparing against 'main'.
+# ./build_tools/pytype/check_diff.sh
+# A specific branch can be specified.
+# ./build_tools/pytype/check_diff.sh google
+# Or all python files outside of './third_party/' can be checked.
+# ./build_tools/pytype/check_diff.sh all
+
+DIFF_TARGET="${1:-main}"
+echo "Running pycheck against '${DIFF_TARGET?}'"
+
+if [[ "${DIFF_TARGET?}" = "all" ]]; then
+ FILES=$(find -name "*\.py" -not -path "./third_party/*")
+else
+ FILES=$(git diff --name-only "${DIFF_TARGET?}" | grep '.*\.py')
+fi
+
+
+# We seperate the python files into multiple pytype calls because otherwise
+# Ninja gets confused. See https://github.com/google/pytype/issues/198
+BASE=$(echo "${FILES?}" | grep -vP '^(\./)?integrations/*')
+IREE_TF=$(echo "${FILES?}" | \
+ grep -P '^(\./)?integrations/tensorflow/bindings/python/pyiree/tf/.*')
+IREE_XLA=$(echo "${FILES?}" | \
+ grep -P '^(\./)?integrations/tensorflow/bindings/python/pyiree/xla/.*')
+COMPILER=$(echo "${FILES?}" | \
+ grep -P '^(\./)?integrations/tensorflow/compiler/.*')
+E2E=$(echo "${FILES?}" | grep -P '^(\./)?integrations/tensorflow/e2e/.*')
+
+function check_files() {
+ # $1: previous return code
+ # $2...: files to check
+ if [[ -z "${@:2}" ]]; then
+ echo "No files to check."
+ echo
+ return "${1?}"
+ fi
+
+ # We disable import-error because pytype doesn't have access to bazel.
+ # We disable pyi-error because of the way the bindings imports work.
+ echo "${@:2}" | \
+ xargs python3 -m pytype --disable=import-error,pyi-error -j $(nproc)
+ EXIT_CODE="$?"
+ echo
+ if [[ "${EXIT_CODE?}" -gt "${1?}" ]]; then
+ return "${EXIT_CODE?}"
+ else
+ return "${1?}"
+ fi
+}
+
+MAX_CODE=0
+
+echo "Checking .py files outside of integrations/"
+check_files "${MAX_CODE?}" "${BASE?}"
+MAX_CODE="$?"
+
+echo "Checking .py files in integrations/tensorflow/bindings/python/pyiree/tf/.*"
+check_files "${MAX_CODE?}" "${IREE_TF?}"
+MAX_CODE="$?"
+
+echo "Checking .py files in integrations/tensorflow/bindings/python/pyiree/xla/.*"
+check_files "${MAX_CODE?}" "${IREE_XLA?}"
+MAX_CODE="$?"
+
+echo "Checking .py files in integrations/tensorflow/compiler/.*"
+check_files "${MAX_CODE?}" "${COMPILER?}"
+MAX_CODE="$?"
+
+echo "Checking .py files in integrations/tensorflow/e2e/.*"
+check_files "${MAX_CODE?}" "${E2E?}"
+MAX_CODE="$?"
+
+
+if [[ "${MAX_CODE?}" -ne "0" ]]; then
+ echo "One or more pytype checks failed."
+ echo "You can view these errors locally by running"
+ echo " ./build_tools/pytype/check_diff.sh ${DIFF_TARGET?}"
+ exit "${MAX_CODE?}"
+fi
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py
index b09f95b..111ad46 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py
+++ b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py
@@ -249,7 +249,7 @@
"""Stores the inputs and outputs of a series of calls to a module."""
def __init__(self,
- module: tf_utils.CompiledModule,
+ module: Union[tf_utils.CompiledModule, None],
function: Union[Callable[["TracedModule"], None], None],
_load_dict: Dict[str, Any] = None):
"""Extracts metadata from module and function and initializes.
@@ -563,7 +563,7 @@
self._module = module
self._trace = trace
- def _trace_call(self, method: Callable[..., Any], method_name: str):
+ def _trace_call(self, method: tf_utils._FunctionWrapper, method_name: str):
"""Decorates a CompiledModule method to capture its inputs and outputs."""
def call(*args, **kwargs):
@@ -611,8 +611,8 @@
def compile_tf_module(
- module_class: Type[tf.Module], exported_names: Sequence[str] = ()
-) -> Callable[[Any], Any]:
+ module_class: Type[tf.Module],
+ exported_names: Sequence[str] = ()) -> Modules:
"""Compiles module_class to each backend that we test.
Args:
@@ -648,11 +648,10 @@
return _global_modules
-def compile_tf_signature_def_saved_model(saved_model_dir: str,
- saved_model_tags: Set[str],
- module_name: str, exported_name: str,
- input_names: Sequence[str],
- output_names: Sequence[str]):
+def compile_tf_signature_def_saved_model(
+ saved_model_dir: str, saved_model_tags: Set[str], module_name: str,
+ exported_name: str, input_names: Sequence[str],
+ output_names: Sequence[str]) -> Modules:
"""Compiles a SignatureDef SavedModel to each backend that we test.
Args:
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
index 99573c5..c0f8473 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
+++ b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
@@ -87,10 +87,10 @@
def _setup_mlir_crash_reproducer(
- function: Callable[[Any], Any],
+ function: Any, # pytype doesn't support arbitrary Callable[*args, **kwargs]
artifacts_dir: str,
backend_id: str,
-) -> Callable[[Any], Any]:
+) -> Any: # Callable[Any, Any]
"""Wraps `function` so that it a MLIR crash reproducer is saved if it crashes.
Writes to `artifacts_dir/reproducer__{backend}.mlir` in the case of a crash.
@@ -253,6 +253,16 @@
exported_name, artifacts_dir)
+class _FunctionWrapper(object):
+
+ def __call__(self, *args, **kwargs):
+ raise NotImplementedError()
+
+ def get_serialized_values(self) -> Tuple[Tuple[str], Tuple[str]]:
+ """Dummy function to match _IreeFunctionWrapper's API."""
+ return ("",), ("",)
+
+
class CompiledModule(object):
"""Base class for the TF and IREE compiled modules."""
@@ -260,7 +270,7 @@
self,
module_name: str,
backend_info: "BackendInfo",
- compiled_paths: Dict[str, str],
+ compiled_paths: Union[Dict[str, str], None],
):
"""Shared base constructor – not useful on its own.
@@ -344,6 +354,9 @@
"""
raise NotImplementedError()
+ def __getattr__(self, attr: str) -> _FunctionWrapper:
+ raise NotImplementedError()
+
def iree_serializable(self):
return False
@@ -351,13 +364,6 @@
return False
-class _FunctionWrapper(object):
-
- def get_serialized_values(self) -> Tuple[Tuple[str], Tuple[str]]:
- """Dummy function to match _IreeFunctionWrapper's API."""
- return (), ()
-
-
class _IreeFunctionWrapper(_FunctionWrapper):
"""Wraps an IREE function, making it callable."""
@@ -681,7 +687,7 @@
instance = module_class()
functions = []
for name in exported_names:
- functions.append(instance.__getattribute__(name).get_concrete_function())
+ functions.append(getattr(instance, name).get_concrete_function())
return functions, exported_names
@@ -787,7 +793,8 @@
def __init__(self, interpreter: tf.lite.Interpreter):
self._interpreter = interpreter
- def __call__(self, *args, **kwargs) -> Tuple[Any]:
+ def __call__(self, *args,
+ **kwargs) -> Union[Dict[str, Any], Tuple[Any], np.ndarray]:
if len(args) and len(kwargs):
raise ValueError("Passing both args and kwargs is not supported by "
"_TfLiteFunctionWrapper")
@@ -823,13 +830,12 @@
outputs.append(value)
# Process them to match the output of the tf.Module.
- if not is_dict:
- outputs = tuple(outputs)
- if len(outputs) == 1:
- outputs = outputs[0]
+ if is_dict:
+ return dict(outputs)
else:
- outputs = dict(outputs)
- return outputs
+ if len(outputs) == 1:
+ return outputs[0]
+ return tuple(outputs)
class TfLiteCompiledModule(CompiledModule):
diff --git a/integrations/tensorflow/compiler/BUILD b/integrations/tensorflow/compiler/BUILD
index 838e202..fba7d6d 100644
--- a/integrations/tensorflow/compiler/BUILD
+++ b/integrations/tensorflow/compiler/BUILD
@@ -21,6 +21,7 @@
cc_library(
name = "tensorflow",
srcs = [
+ "CheckNoTF.cpp",
"LegalizeTF.cpp",
"Passes.cpp",
"PropagateResourceCasts.cpp",
diff --git a/integrations/tensorflow/compiler/CheckNoTF.cpp b/integrations/tensorflow/compiler/CheckNoTF.cpp
new file mode 100644
index 0000000..246f28a
--- /dev/null
+++ b/integrations/tensorflow/compiler/CheckNoTF.cpp
@@ -0,0 +1,93 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
+#include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h"
+#include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/lower_tf.h"
+#include "tensorflow/compiler/mlir/xla/transforms/passes.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace {
+
+class CheckNoTensorflow : public PassWrapper<CheckNoTensorflow, FunctionPass> {
+ public:
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<chlo::HloClientDialect, mhlo::MhloDialect,
+ shape::ShapeDialect, StandardOpsDialect>();
+ }
+
+ CheckNoTensorflow() = default;
+ CheckNoTensorflow(const CheckNoTensorflow &) {}
+
+ /// Validates that no TensorFlow frontends ops are in the function.
+ void runOnFunction() override {
+ auto op = getFunction();
+ auto context = op.getContext();
+
+ Dialect *dialect = context->getLoadedDialect("tf");
+ DenseSet<Operation *> illegalOps;
+ op.walk([&](Operation *op) {
+ if (op->getDialect() == dialect) {
+ illegalOps.insert(op);
+ }
+ });
+
+ if (!illegalOps.empty()) {
+ emitLegalizationErrors(op, illegalOps);
+ return signalPassFailure();
+ }
+ }
+
+ // Emits debug information which includes the number of ops of each type which
+ // failed to legalize.
+ void emitLegalizationErrors(Operation *op,
+ const DenseSet<Operation *> &nonlegalizedOps) {
+ // Print op errors for each of the TensorFlow ops that still remain.
+ std::map<StringRef, int> opNameCounts;
+ for (Operation *nonlegalizedOp : nonlegalizedOps) {
+ StringRef opName = nonlegalizedOp->getName().getStringRef();
+ opNameCounts[opName]++;
+ nonlegalizedOp->emitOpError()
+ << ": unlegalized TensorFlow op still exists";
+ }
+
+ std::vector<std::string> errorMessages;
+ errorMessages.reserve(opNameCounts.size());
+ for (const auto &opInfo : opNameCounts) {
+ errorMessages.push_back(
+ llvm::formatv("\t{0} (count: {1})", opInfo.first, opInfo.second));
+ }
+ Location loc = op->getLoc();
+ emitError(loc) << "The following Tensorflow operations still remain: \n"
+ << llvm::join(errorMessages, "\n") << "\n";
+ }
+};
+
+static PassRegistration<CheckNoTensorflow> pass(
+ "iree-check-no-tf", "Check that no TensorFlow frontend ops remain");
+} // namespace
+
+std::unique_ptr<OperationPass<FuncOp>> createCheckNoTF() {
+ return std::make_unique<CheckNoTensorflow>();
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/integrations/tensorflow/compiler/Passes.cpp b/integrations/tensorflow/compiler/Passes.cpp
index b4b6c47..db2e9d9 100644
--- a/integrations/tensorflow/compiler/Passes.cpp
+++ b/integrations/tensorflow/compiler/Passes.cpp
@@ -64,6 +64,11 @@
// - It removes tf_saved_model.semantics from the module, which we can only
// do at the very end.
pm.addPass(createTFSavedModelLowerExportedFunctions());
+
+ ////////////////////////////////////////////////////////////////////////////
+ // Validate that all Tensorflow has been legalized away.
+ ////////////////////////////////////////////////////////////////////////////
+ pm.addPass(createCheckNoTF());
}
static mlir::PassPipelineRegistration<> pipeline(
diff --git a/integrations/tensorflow/compiler/Passes.h b/integrations/tensorflow/compiler/Passes.h
index e3016a3..8b3b219 100644
--- a/integrations/tensorflow/compiler/Passes.h
+++ b/integrations/tensorflow/compiler/Passes.h
@@ -40,6 +40,9 @@
// Push resource casts forward to better propagate resource related shapes.
std::unique_ptr<OperationPass<ModuleOp>> createPropagateResourceCasts();
+// Validates whether any Tensorflow operations remain.
+std::unique_ptr<OperationPass<FuncOp>> createCheckNoTF();
+
// Create a single pipeline that will run all the needed IREE-specific TF import
// passes in the right order.
void createIreeTfImportPipeline(OpPassManager &pm);
diff --git a/integrations/tensorflow/compiler/test/check-no-tf.mlir b/integrations/tensorflow/compiler/test/check-no-tf.mlir
new file mode 100644
index 0000000..a7c49e6
--- /dev/null
+++ b/integrations/tensorflow/compiler/test/check-no-tf.mlir
@@ -0,0 +1,28 @@
+// RUN: iree-tf-opt %s -iree-check-no-tf -split-input-file -verify-diagnostics
+
+// CHECK-LABEL: func @f
+func @f() -> (tensor<i32>) {
+ // CHECK: [[VAL0:%.+]] = mhlo.constant dense<3>
+ %0 = mhlo.constant dense<3> : tensor<i32>
+ return %0 : tensor<i32>
+}
+
+// -----
+
+// expected-error@+3 {{'tf.Const' op : unlegalized TensorFlow op still exists}}
+// expected-error@below {{The following Tensorflow operations still remain}}
+func @f() -> (tensor<i32>) {
+ %0 = "tf.Const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+ return %0 : tensor<i32>
+}
+
+// -----
+
+// expected-error@+4 {{'tf.Const' op : unlegalized TensorFlow op still exists}}
+// expected-error@+4 {{'tf.Add' op : unlegalized TensorFlow op still exists}}
+// expected-error@below {{The following Tensorflow operations still remain}}
+func @f(%arg0 : tensor<i32>) -> (tensor<i32>) {
+ %0 = "tf.Const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+ %1 = "tf.Add"(%arg0, %0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+ return %1 : tensor<i32>
+}
diff --git a/integrations/tensorflow/compiler/test/saved_model_adopt_exports.py b/integrations/tensorflow/compiler/test/saved_model_adopt_exports.py
index 256ce89..aff5dd8 100644
--- a/integrations/tensorflow/compiler/test/saved_model_adopt_exports.py
+++ b/integrations/tensorflow/compiler/test/saved_model_adopt_exports.py
@@ -23,6 +23,7 @@
SAVED_MODEL_IMPORT_PASSES = [
"tf-executor-graph-pruning",
"tf-standard-pipeline",
+ "iree-xla-legalize-tf",
"iree-tf-import-pipeline",
"canonicalize",
]
@@ -114,8 +115,8 @@
# CHECK: attributes
# CHECK-SAME: iree.module.export
# CHECK-SAME: iree.reflection = {abi = "sip", abiv = 1 : i32, sip = "I1!R1!"}
-# CHECK-DAG: [[CONST_2xf32:%.+]] = "tf.Const"() {value = dense<[0.000000e+00, 1.000000e+00]> : tensor<2xf32>} : () -> tensor<2xf32>
-# CHECK-DAG: [[CONST_3xf32:%.+]] = "tf.Const"() {value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00]> : tensor<3xf32>} : () -> tensor<3xf32>
+# CHECK-DAG: [[CONST_2xf32:%.+]] = mhlo.constant dense<[0.000000e+00, 1.000000e+00]>
+# CHECK-DAG: [[CONST_3xf32:%.+]] = mhlo.constant dense<[0.000000e+00, 1.000000e+00, 2.000000e+00]>
# CHECK-DAG: flow.variable.store [[CONST_2xf32]], @v : tensor<2xf32>
# CHECK-DAG: flow.variable.store [[CONST_3xf32]], @v : tensor<3xf32>
# CHECK: FINISH_TEST
diff --git a/integrations/tensorflow/e2e/BUILD b/integrations/tensorflow/e2e/BUILD
index bc8a998..e919d3e 100644
--- a/integrations/tensorflow/e2e/BUILD
+++ b/integrations/tensorflow/e2e/BUILD
@@ -201,31 +201,11 @@
"tflite": ["mobile_bert_squad_test.py"],
"iree_vmla": ["mobile_bert_squad_test.py"],
"iree_llvmjit": ["mobile_bert_squad_test.py"],
- },
- reference_backend = "tf",
- tags = [
- "external",
- "guitar",
- "manual",
- "no-remote",
- "nokokoro",
- "notap",
- ],
- deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
- "//integrations/tensorflow/bindings/python/pyiree/tf/support",
- ],
-)
-
-iree_e2e_test_suite(
- name = "mobile_bert_squad_tests_failing",
- size = "enormous",
- backends_to_srcs = {
"iree_vulkan": ["mobile_bert_squad_test.py"],
},
reference_backend = "tf",
tags = [
"external",
- "failing",
"guitar",
"manual",
"no-remote",
diff --git a/iree/base/BUILD b/iree/base/BUILD
index a1b46ba..1cc8c50 100644
--- a/iree/base/BUILD
+++ b/iree/base/BUILD
@@ -295,8 +295,10 @@
srcs = ["logging.cc"],
hdrs = ["logging.h"],
deps = [
+ ":tracing",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/flags:flag",
+ "@com_google_absl//absl/strings:str_format",
],
)
diff --git a/iree/base/CMakeLists.txt b/iree/base/CMakeLists.txt
index 257618b..e6ec69a 100644
--- a/iree/base/CMakeLists.txt
+++ b/iree/base/CMakeLists.txt
@@ -345,8 +345,10 @@
SRCS
"logging.cc"
DEPS
+ ::tracing
absl::core_headers
absl::flags
+ absl::str_format
PUBLIC
)
diff --git a/iree/base/logging.cc b/iree/base/logging.cc
index 51ce60b..3b31a9a 100644
--- a/iree/base/logging.cc
+++ b/iree/base/logging.cc
@@ -17,6 +17,8 @@
#include <string>
#include "absl/flags/flag.h"
+#include "absl/strings/str_format.h"
+#include "iree/base/tracing.h"
ABSL_FLAG(int, iree_minloglevel, 0,
"Minimum logging level. 0 = INFO and above.");
@@ -88,6 +90,19 @@
// TODO(scotttodd): Include current system time
fprintf(stderr, "%c %s:%d] %s\n", "IWEF"[severity_], file_name_, line_,
str().c_str());
+
+#if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_LOG_MESSAGES
+ constexpr int kLevelColors[4] = {
+ IREE_TRACING_MESSAGE_LEVEL_INFO, // INFO
+ IREE_TRACING_MESSAGE_LEVEL_WARNING, // WARNING
+ IREE_TRACING_MESSAGE_LEVEL_ERROR, // ERROR
+ IREE_TRACING_MESSAGE_LEVEL_ERROR, // FATAL
+ };
+ std::string message =
+ absl::StrFormat("%s:%d] %s\n", file_name_, line_, str().c_str());
+ IREE_TRACE_MESSAGE_DYNAMIC_COLORED(kLevelColors[severity_], message.c_str(),
+ message.size());
+#endif // IREE_TRACING_FEATURES& IREE_TRACING_FEATURE_LOG_MESSAGES
}
LogMessageFatal::LogMessageFatal(const char* file, int line)
diff --git a/iree/base/tracing.h b/iree/base/tracing.h
index a3226ad..98db68e 100644
--- a/iree/base/tracing.h
+++ b/iree/base/tracing.h
@@ -77,6 +77,10 @@
// Tracy UI.
#define IREE_TRACING_FEATURE_SLOW_LOCKS (1 << 5)
+// Forwards log messages to traces, which will be visible under "Messages" in
+// the Tracy UI.
+#define IREE_TRACING_FEATURE_LOG_MESSAGES (1 << 6)
+
#if !defined(IREE_TRACING_MAX_CALLSTACK_DEPTH)
// Tracing functions that capture stack traces will only capture up to N frames.
// The overhead for stack walking scales linearly with the number of frames
@@ -95,7 +99,7 @@
// overridden it with more specific settings.
//
// IREE_TRACING_MODE = 0: tracing disabled
-// IREE_TRACING_MODE = 1: instrumentation and basic statistics
+// IREE_TRACING_MODE = 1: instrumentation, log messages, and basic statistics
// IREE_TRACING_MODE = 2: same as 1 with added allocation tracking
// IREE_TRACING_MODE = 3: same as 2 with callstacks for allocations
// IREE_TRACING_MODE = 4: same as 3 with callstacks for all instrumentation
@@ -105,20 +109,23 @@
#undef IREE_TRACING_MAX_CALLSTACK_DEPTH
#define IREE_TRACING_MAX_CALLSTACK_DEPTH 0
#elif defined(IREE_TRACING_MODE) && IREE_TRACING_MODE == 2
-#define IREE_TRACING_FEATURES \
- (IREE_TRACING_FEATURE_INSTRUMENTATION | \
- IREE_TRACING_FEATURE_ALLOCATION_TRACKING)
-#elif defined(IREE_TRACING_MODE) && IREE_TRACING_MODE == 3
#define IREE_TRACING_FEATURES \
(IREE_TRACING_FEATURE_INSTRUMENTATION | \
IREE_TRACING_FEATURE_ALLOCATION_TRACKING | \
- IREE_TRACING_FEATURE_ALLOCATION_CALLSTACKS)
+ IREE_TRACING_FEATURE_LOG_MESSAGES)
+#elif defined(IREE_TRACING_MODE) && IREE_TRACING_MODE == 3
+#define IREE_TRACING_FEATURES \
+ (IREE_TRACING_FEATURE_INSTRUMENTATION | \
+ IREE_TRACING_FEATURE_ALLOCATION_TRACKING | \
+ IREE_TRACING_FEATURE_ALLOCATION_CALLSTACKS | \
+ IREE_TRACING_FEATURE_LOG_MESSAGES)
#elif defined(IREE_TRACING_MODE) && IREE_TRACING_MODE >= 4
#define IREE_TRACING_FEATURES \
(IREE_TRACING_FEATURE_INSTRUMENTATION | \
IREE_TRACING_FEATURE_INSTRUMENTATION_CALLSTACKS | \
IREE_TRACING_FEATURE_ALLOCATION_TRACKING | \
- IREE_TRACING_FEATURE_ALLOCATION_CALLSTACKS)
+ IREE_TRACING_FEATURE_ALLOCATION_CALLSTACKS | \
+ IREE_TRACING_FEATURE_LOG_MESSAGES)
#else
#define IREE_TRACING_FEATURES 0
#endif // IREE_TRACING_MODE
@@ -336,11 +343,21 @@
// The message text must be a compile-time string literal.
#define IREE_TRACE_MESSAGE(level, value_literal) \
___tracy_emit_messageLC(value_literal, IREE_TRACING_MESSAGE_LEVEL_##level, 0)
+// Logs a message with the given color to the trace.
+// Standard colors are defined as IREE_TRACING_MESSAGE_LEVEL_* values.
+// The message text must be a compile-time string literal.
+#define IREE_TRACE_MESSAGE_COLORED(color, value_literal) \
+ ___tracy_emit_messageLC(value_literal, color, 0)
// Logs a dynamically-allocated message at the given logging level to the trace.
// The string |value| will be copied into the trace buffer.
#define IREE_TRACE_MESSAGE_DYNAMIC(level, value, value_length) \
___tracy_emit_messageC(value, value_length, \
IREE_TRACING_MESSAGE_LEVEL_##level, 0)
+// Logs a dynamically-allocated message with the given color to the trace.
+// Standard colors are defined as IREE_TRACING_MESSAGE_LEVEL_* values.
+// The string |value| will be copied into the trace buffer.
+#define IREE_TRACE_MESSAGE_DYNAMIC_COLORED(color, value, value_length) \
+ ___tracy_emit_messageC(value, value_length, color, 0)
// Utilities:
#define IREE_TRACE_IMPL_GET_VARIADIC_HELPER_(_1, _2, _3, NAME, ...) NAME
@@ -370,7 +387,9 @@
#define IREE_TRACE_FRAME_MARK_BEGIN_NAMED(name_literal)
#define IREE_TRACE_FRAME_MARK_END_NAMED(name_literal)
#define IREE_TRACE_MESSAGE(level, value_literal)
+#define IREE_TRACE_MESSAGE_COLORED(color, value_literal)
#define IREE_TRACE_MESSAGE_DYNAMIC(level, value, value_length)
+#define IREE_TRACE_MESSAGE_DYNAMIC_COLORED(color, value, value_length)
#endif // IREE_TRACING_FEATURE_INSTRUMENTATION
//===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/ConvertToSPIRVPass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/ConvertToSPIRVPass.cpp
index 7a788b1..93b8809 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/ConvertToSPIRVPass.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/ConvertToSPIRVPass.cpp
@@ -130,11 +130,8 @@
spirv::GlobalVariableOp getOrInsertResourceVariable(Location loc, Type type,
unsigned set,
unsigned binding,
- StringRef funcName,
Block &block) {
- auto name =
- llvm::formatv("__resource_var_{0}_{1}_{2}__", set, binding, funcName)
- .str();
+ auto name = llvm::formatv("__resource_var_{0}_{1}__", set, binding).str();
for (auto varOp : block.getOps<spirv::GlobalVariableOp>()) {
if (varOp.sym_name() == name) return varOp;
}
@@ -371,10 +368,9 @@
SymbolTable::lookupNearestSymbolFrom(
phOp, phOp.getAttrOfType<SymbolRefAttr>("binding")));
- StringRef funcName = phOp.getParentOfType<spirv::FuncOp>().getName();
- spirv::GlobalVariableOp varOp = getOrInsertResourceVariable(
- phOp.getLoc(), convertedType, bindingOp.set(), bindingOp.binding(),
- funcName, *moduleOp.getBody());
+ spirv::GlobalVariableOp varOp =
+ getOrInsertResourceVariable(phOp.getLoc(), convertedType, bindingOp.set(),
+ bindingOp.binding(), *moduleOp.getBody());
rewriter.replaceOpWithNewOp<spirv::AddressOfOp>(phOp, varOp);
return success();
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/Passes.cpp b/iree/compiler/Conversion/LinalgToSPIRV/Passes.cpp
index bb6faf8..d751065 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/Passes.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/Passes.cpp
@@ -63,6 +63,11 @@
llvm::cl::desc(
"Enable use of vectorization in SPIR-V code generation pipeline"),
llvm::cl::init(false)};
+ Option<bool> useVectorPass{
+ *this, "use-vector-pass",
+ llvm::cl::desc("Enable use of Linalg vectorization in SPIR-V code "
+ "generation pipeline"),
+ llvm::cl::init(false)};
};
static void addLinalgToSPIRVPasses(OpPassManager &pm,
@@ -94,7 +99,9 @@
//===--------------------------------------------------------------------===//
pm.addPass(createSplitDispatchFunctionPass());
pm.addPass(createLinalgTileAndFusePass(options));
- pm.addPass(createLoadStoreVectorizationPass());
+ if (options.useVectorPass) {
+ pm.addPass(createLoadStoreVectorizationPass());
+ }
pm.addPass(createCanonicalizerPass());
//===--------------------------------------------------------------------===//
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/Passes.h b/iree/compiler/Conversion/LinalgToSPIRV/Passes.h
index 502db6a..80805e2 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/Passes.h
+++ b/iree/compiler/Conversion/LinalgToSPIRV/Passes.h
@@ -29,6 +29,7 @@
SmallVector<int64_t, 3> tileSizes = {};
bool useWorkgroupMemory = false;
bool useVectorization = false;
+ bool useVectorPass = false;
};
/// Pass to initialize the function that computes the number of workgroups for
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/VectorToGPUPass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/VectorToGPUPass.cpp
index 7721f41..eb2a32e 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/VectorToGPUPass.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/VectorToGPUPass.cpp
@@ -64,6 +64,7 @@
private:
void tileAndVectorizeLinalgCopy(FuncOp funcOp, MLIRContext *context);
+ void lowerVectorOps(FuncOp funcOp, MLIRContext *context);
};
// Common class for all vector to GPU patterns.
@@ -149,8 +150,7 @@
loc, rewriter.getIndexType(), rewriter.getStringAttr("x"));
Value index = rewriter.create<AddIOp>(loc, ThreadIndex, indices.back());
indices.back() = index;
- rewriter.create<StoreOp>(op.getLoc(), operands[0], operands[1], indices);
- rewriter.eraseOp(op);
+ rewriter.replaceOpWithNewOp<StoreOp>(op, operands[0], operands[1], indices);
return success();
}
};
@@ -204,11 +204,141 @@
applyPatternsAndFoldGreedily(funcOp, vectorizationPatterns);
}
+// Convert vector transfer_read to a load if possible. This is the case only if
+// the element type of the memref matches the element type we want to load.
+class VectorTransferReadToLoad
+ : public OpRewritePattern<vector::TransferReadOp> {
+ public:
+ using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::TransferReadOp op,
+ PatternRewriter &rewriter) const override {
+ if (op.getVectorType().getNumElements() != 1 ||
+ op.getMemRefType().getElementType() !=
+ op.getVectorType().getElementType()) {
+ return failure();
+ }
+ auto loc = op.getLoc();
+ Value newOp = rewriter.create<LoadOp>(loc, op.memref(), op.indices());
+ newOp =
+ rewriter.create<vector::BroadcastOp>(loc, op.getVectorType(), newOp);
+ rewriter.replaceOp(op, newOp);
+ return success();
+ }
+};
+
+// Convert vector transfer_write to a store if possible. This is the case only
+// if the element type of the memref matches the element type we want to store.
+class VectorTransferWriteToStore
+ : public OpRewritePattern<vector::TransferWriteOp> {
+ public:
+ using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::TransferWriteOp op,
+ PatternRewriter &rewriter) const override {
+ if (op.getVectorType().getNumElements() != 1 ||
+ op.getMemRefType().getElementType() !=
+ op.getVectorType().getElementType()) {
+ return failure();
+ }
+ auto loc = op.getLoc();
+ SmallVector<int64_t, 2> zero(op.getVectorType().getRank(), 0);
+ Value scalarValue =
+ rewriter.create<vector::ExtractOp>(loc, op.vector(), zero);
+ rewriter.create<StoreOp>(loc, scalarValue, op.memref(), op.indices());
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+
+// Lower vector contract to a single scalar or vector mulf+addf. Insert casts to
+// convert from 2D vector to 1D vector or scalar.
+class VectorContractLowering : public OpRewritePattern<vector::ContractionOp> {
+ public:
+ using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::ContractionOp op,
+ PatternRewriter &rewriter) const override {
+ auto iteratorTypes = op.iterator_types().getValue();
+ if (iteratorTypes.size() != 3 || !isParallelIterator(iteratorTypes[0]) ||
+ !isParallelIterator(iteratorTypes[1]) ||
+ !isReductionIterator(iteratorTypes[2]) ||
+ !isRowMajorMatmul(op.indexing_maps())) {
+ return failure();
+ }
+ if (op.getLhsType().getNumElements() != 1) return failure();
+ unsigned vecSize = op.getAccType().cast<VectorType>().getNumElements();
+ if (!(vecSize >= 1 && vecSize <= 4)) return failure();
+ auto loc = op.getLoc();
+ VectorType vecType = VectorType::get(
+ vecSize, op.getResultType().cast<VectorType>().getElementType());
+ std::array<int64_t, 2> zero = {0, 0};
+ Value lhs = rewriter.create<vector::ExtractOp>(loc, op.lhs(), zero);
+ Value rhs, acc;
+ if (vecSize == 1) {
+ rhs = rewriter.create<vector::ExtractOp>(loc, op.rhs(), zero);
+ acc = rewriter.create<vector::ExtractOp>(loc, op.acc(), zero);
+ } else {
+ lhs = rewriter.create<vector::BroadcastOp>(loc, vecType, lhs);
+ rhs = rewriter.create<vector::ShapeCastOp>(loc, vecType, op.rhs());
+ acc = rewriter.create<vector::ShapeCastOp>(loc, vecType, op.acc());
+ }
+ Value newOp = rewriter.create<MulFOp>(loc, lhs, rhs);
+ newOp = rewriter.create<AddFOp>(loc, newOp, acc);
+ if (vecSize == 1)
+ newOp =
+ rewriter.create<vector::BroadcastOp>(loc, op.getResultType(), newOp);
+ else
+ newOp =
+ rewriter.create<vector::ShapeCastOp>(loc, op.getResultType(), newOp);
+ rewriter.replaceOp(op, newOp);
+ 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.
+class ExtractStridedLowering
+ : public OpRewritePattern<vector::ExtractStridedSliceOp> {
+ public:
+ using OpRewritePattern<vector::ExtractStridedSliceOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::ExtractStridedSliceOp op,
+ PatternRewriter &rewriter) const override {
+ // Only handle cases extracting a degenerated vector so that we can generate
+ // an extractOp with scalar destination.
+ if (op.getResult().getType().cast<VectorType>().getNumElements() != 1)
+ return failure();
+ auto loc = op.getLoc();
+ SmallVector<int64_t, 4> offsets = llvm::to_vector<4>(
+ llvm::map_range(op.offsets().getAsRange<IntegerAttr>(),
+ [](IntegerAttr attr) { return attr.getInt(); }));
+ offsets.resize(op.getVectorType().getRank(), 0);
+ Value newOp = rewriter.create<vector::ExtractOp>(loc, op.vector(), offsets);
+ newOp = rewriter.create<vector::BroadcastOp>(loc, op.getResult().getType(),
+ newOp);
+ rewriter.replaceOp(op, newOp);
+ return success();
+ }
+};
+
+// Lower vector ops to instructions that can be later converted to SPIR-V.
+void ConvertVectorToGPUPass::lowerVectorOps(FuncOp funcOp,
+ MLIRContext *context) {
+ OwningRewritePatternList patterns;
+ patterns.insert<VectorContractLowering, VectorTransferReadToLoad,
+ VectorTransferWriteToStore, ExtractStridedLowering>(context);
+ applyPatternsAndFoldGreedily(funcOp, patterns);
+}
+
void ConvertVectorToGPUPass::runOnOperation() {
MLIRContext *context = &getContext();
FuncOp funcOp = getOperation();
tileAndVectorizeLinalgCopy(funcOp, context);
+ lowerVectorOps(funcOp, context);
+
auto &cooperativeMatrixAnalysis = getAnalysis<CooperativeMatrixAnalysis>();
OwningRewritePatternList patterns;
patterns.insert<UnaryAndBinaryOpPattern<AddFOp>, VectorTransferReadConversion,
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_spirv.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_spirv.mlir
index f6bbc40..ac83d70 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_spirv.mlir
+++ b/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_spirv.mlir
@@ -23,12 +23,12 @@
module attributes {spv.target_env = #spv.target_env<#spv.vce<v1.3, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
- // CHECK: spv.globalVariable @__resource_var_3_4_resource_variable__ bind(3, 4) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer>
- // CHECK: spv.globalVariable @__resource_var_1_2_resource_variable__ bind(1, 2) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer>
+ // CHECK: spv.globalVariable @__resource_var_3_4__ bind(3, 4) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer>
+ // CHECK: spv.globalVariable @__resource_var_1_2__ bind(1, 2) : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer>
// CHECK: spv.func @resource_variable()
func @resource_variable() {
- // CHECK: spv._address_of @__resource_var_1_2_resource_variable__ : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer>
- // CHECK: spv._address_of @__resource_var_3_4_resource_variable__ : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer>
+ // CHECK: spv._address_of @__resource_var_1_2__ : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer>
+ // CHECK: spv._address_of @__resource_var_3_4__ : !spv.ptr<!spv.struct<(!spv.array<16 x f32, stride=4> [0])>, StorageBuffer>
%0 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0} : memref<4x4xf32>
%1 = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<4x4xf32>
return
@@ -42,24 +42,6 @@
// -----
-module attributes {spv.target_env = #spv.target_env<#spv.vce<v1.3, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>, vkspv.entry_point_schedule = ["ForwardPass_ex_dispatch_8_dispatch_0", "ForwardPass_ex_dispatch_8_dispatch_1"]} {
- // CHECK: spv.globalVariable @__resource_var_0_0_dispatch_0__ bind(0, 0) : !spv.ptr<!spv.struct<(!spv.array<1 x vector<4xf32>, stride=16> [0])>, StorageBuffer>
- // CHECK: spv.globalVariable @__resource_var_0_0_dispatch_1__ bind(0, 0) : !spv.ptr<!spv.struct<(!spv.array<4 x f32, stride=4> [0])>, StorageBuffer>
- func @dispatch_1() attributes {spv.entry_point_abi = {local_size = dense<[32, 1, 1]> : vector<3xi32>}, vkspv.num_workgroups_fn = @ForwardPass_ex_dispatch_8_dispatch_1__num_workgroups__} {
- %0 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<4xf32>
- return
- }
- func @dispatch_0() attributes {spv.entry_point_abi = {local_size = dense<[32, 1, 1]> : vector<3xi32>}, vkspv.num_workgroups_fn = @ForwardPass_ex_dispatch_8_dispatch_0__num_workgroups__} {
- %0 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<1xvector<4xf32>>
- return
- }
- hal.interface @legacy_io attributes {sym_visibility = "private"} {
- hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
- }
-}
-
-// -----
-
#map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
#map2 = affine_map<(d0, d1, d2) -> (d2, d1)>
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/vector_to_gpu.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/vector_to_gpu.mlir
index 30d94c5..99ec424 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/test/vector_to_gpu.mlir
+++ b/iree/compiler/Conversion/LinalgToSPIRV/test/vector_to_gpu.mlir
@@ -59,3 +59,75 @@
// CHECK: %[[LOAD:.+]] = vector.transfer_read %[[SVs]][%c0, %c0], %cst {{.*}} : memref<1x4xf32, {{.*}}>, vector<1x4xf32>
// CHECK: vector.transfer_write %[[LOAD]], %[[SVd]][%[[C0]], %[[C0]]] {{.*}} : vector<1x4xf32>, memref<1x4xf32
}
+
+// -----
+
+module attributes {gpu.container_module, spv.target_env = #spv.target_env<#spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
+ func @transfer_ops(%arg0: memref<32x32xf32>, %arg1 : vector<1x1xf32>) -> vector<1x1xf32> attributes {spv.entry_point_abi = {local_size = dense<[128, 1, 1]> : vector<3xi32>}} {
+ %c0 = constant 0 : index
+ %cst = constant 0.0 : f32
+ %0 = vector.transfer_read %arg0[%c0, %c0], %cst : memref<32x32xf32>, vector<1x1xf32>
+ vector.transfer_write %arg1, %arg0[%c0, %c0] : vector<1x1xf32>, memref<32x32xf32>
+ return %0 : vector<1x1xf32>
+ }
+ // CHECK-LABEL: func @transfer_ops
+ // CHECK-SAME: (%[[ARG0:.*]]: memref<32x32xf32>, %[[ARG1:.*]]: vector<1x1xf32>
+ // CHECK: %[[C0:.*]] = constant 0 : index
+ // CHECK: %[[LOAD:.*]] = load %[[ARG0]][%[[C0]], %[[C0]]] : memref<32x32xf32>
+ // CHECK: %[[B:.*]] = vector.broadcast %[[LOAD]] : f32 to vector<1x1xf32>
+ // CHECK: %[[EXT:.*]] = vector.extract %[[ARG1]][0, 0] : vector<1x1xf32>
+ // CHECK: store %[[EXT]], %[[ARG0]][%[[C0]], %[[C0]]] : memref<32x32xf32>
+ // CHECK: return %[[B]] : vector<1x1xf32>
+}
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+module attributes {gpu.container_module, spv.target_env = #spv.target_env<#spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
+ func @contract_ops(%arg0 : vector<1x1xf32>, %arg1 : vector<1x4xf32>,
+ %arg2 : vector<1x4xf32>, %arg3 : vector<1x1xf32>,
+ %arg4 : vector<1x1xf32>) -> (vector<1x1xf32>, vector<1x4xf32>) attributes {spv.entry_point_abi = {local_size = dense<[128, 1, 1]> : vector<3xi32>}} {
+ %0 = vector.contract {indexing_maps = [#map0, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"]} %arg0, %arg3, %arg4
+ : vector<1x1xf32>, vector<1x1xf32> into vector<1x1xf32>
+ %1 = vector.contract {indexing_maps = [#map0, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"]} %arg0, %arg1, %arg2
+ : vector<1x1xf32>, vector<1x4xf32> into vector<1x4xf32>
+ return %0, %1 : vector<1x1xf32>, vector<1x4xf32>
+ }
+ // CHECK-LABEL: func @contract_ops
+ // CHECK-SAME: (%[[ARG0:.*]]: vector<1x1xf32>, %[[ARG1:.*]]: vector<1x4xf32>, %[[ARG2:.*]]: vector<1x4xf32>, %[[ARG3:.*]]: vector<1x1xf32>, %[[ARG4:.*]]: vector<1x1xf32>)
+ // CHECK: %[[A:.*]] = vector.extract %[[ARG0]][0, 0] : vector<1x1xf32>
+ // CHECK: %[[B:.*]] = vector.extract %[[ARG3]][0, 0] : vector<1x1xf32>
+ // CHECK: %[[C:.*]] = vector.extract %[[ARG4]][0, 0] : vector<1x1xf32>
+ // CHECK: %[[MUL:.*]] = mulf %[[A]], %[[B]] : f32
+ // CHECK: %[[ADD:.*]] = addf %[[MUL]], %[[C]] : f32
+ // CHECK: %[[R0:.*]] = vector.broadcast %[[ADD]] : f32 to vector<1x1xf32>
+ // CHECK: %[[A:.*]] = vector.extract %[[ARG0]][0, 0] : vector<1x1xf32>
+ // CHECK: %[[VA:.*]] = vector.broadcast %[[A]] : f32 to vector<4xf32>
+ // CHECK: %[[VB:.*]] = vector.shape_cast %[[ARG1]] : vector<1x4xf32> to vector<4xf32>
+ // CHECK: %[[VC:.*]] = vector.shape_cast %[[ARG2]] : vector<1x4xf32> to vector<4xf32>
+ // CHECK: %[[VMUL:.*]] = mulf %[[VA]], %[[VB]] : vector<4xf32>
+ // CHECK: %[[VADD:.*]] = addf %[[VMUL]], %[[VC]] : vector<4xf32>
+ // CHECK: %[[R1:.*]] = vector.shape_cast %[[VADD]] : vector<4xf32> to vector<1x4xf32>
+ // CHECK: return %[[R0]], %[[R1]] : vector<1x1xf32>, vector<1x4xf32>
+}
+
+// -----
+
+module attributes {gpu.container_module, spv.target_env = #spv.target_env<#spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
+ func @extract(%arg0 : vector<1x4xf32>) -> vector<1x1xf32> attributes {spv.entry_point_abi = {local_size = dense<[128, 1, 1]> : vector<3xi32>}} {
+ %0 = vector.extract_strided_slice %arg0
+ {offsets = [0, 2], sizes = [1, 1], strides = [1, 1]}
+ : vector<1x4xf32> to vector<1x1xf32>
+ return %0 : vector<1x1xf32>
+ }
+ // CHECK-LABEL: func @extract
+ // CHECK-SAME: (%[[ARG0:.*]]: vector<1x4xf32>
+ // CHECK: %[[A:.*]] = vector.extract %[[ARG0]][0, 2] : vector<1x4xf32>
+ // CHECK: %[[B:.*]] = vector.broadcast %[[A]] : f32 to vector<1x1xf32>
+ // CHECK: return %[[B]] : vector<1x1xf32>
+}
diff --git a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
index ea4c70c..91de450c 100644
--- a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
@@ -48,6 +48,12 @@
// llvm::cl::OptionCategory halVulkanSPIRVOptionsCategory(
// "IREE Vulkan/SPIR-V backend options");
+ static llvm::cl::opt<bool> clUseVectorPass(
+ "iree-spirv-use-vector-pass",
+ llvm::cl::desc(
+ "Enable use of Linalg vectorization in SPIR-V code generation"),
+ llvm::cl::init(false));
+
static llvm::cl::opt<bool> clUseWorkgroupMemory(
"iree-spirv-use-workgroup-memory",
llvm::cl::desc(
@@ -81,6 +87,7 @@
targetOptions.codegenOptions.tileSizes.assign(clTileSizes.begin(),
clTileSizes.end());
targetOptions.codegenOptions.useWorkgroupMemory = clUseWorkgroupMemory;
+ targetOptions.codegenOptions.useVectorPass = clUseVectorPass;
if (!clVulkanTargetEnv.empty()) {
targetOptions.vulkanTargetEnv = clVulkanTargetEnv;
} else {
diff --git a/iree/test/e2e/vulkan_specific/BUILD b/iree/test/e2e/vulkan_specific/BUILD
index 9686098..f4741ac 100644
--- a/iree/test/e2e/vulkan_specific/BUILD
+++ b/iree/test/e2e/vulkan_specific/BUILD
@@ -45,3 +45,15 @@
driver = "vulkan",
target_backend = "vulkan-spirv",
)
+
+iree_check_single_backend_test_suite(
+ name = "check_vulkan-spirv_vulkan_vector",
+ srcs = [
+ "compare.mlir",
+ "log_plus_one.mlir",
+ "pw_add_multiwg.mlir",
+ ],
+ compiler_flags = ["-iree-spirv-use-vector-pass"],
+ driver = "vulkan",
+ target_backend = "vulkan-spirv",
+)
diff --git a/iree/test/e2e/vulkan_specific/CMakeLists.txt b/iree/test/e2e/vulkan_specific/CMakeLists.txt
index 4260b07..d5bd481 100644
--- a/iree/test/e2e/vulkan_specific/CMakeLists.txt
+++ b/iree/test/e2e/vulkan_specific/CMakeLists.txt
@@ -41,3 +41,18 @@
COMPILER_FLAGS
"-iree-spirv-use-workgroup-memory"
)
+
+iree_check_single_backend_test_suite(
+ NAME
+ check_vulkan-spirv_vulkan_vector
+ SRCS
+ "compare.mlir"
+ "log_plus_one.mlir"
+ "pw_add_multiwg.mlir"
+ TARGET_BACKEND
+ vulkan-spirv
+ DRIVER
+ vulkan
+ COMPILER_FLAGS
+ "-iree-spirv-use-vector-pass"
+)
diff --git a/iree/test/e2e/vulkan_specific/compare.mlir b/iree/test/e2e/vulkan_specific/compare.mlir
new file mode 100644
index 0000000..099670a
--- /dev/null
+++ b/iree/test/e2e/vulkan_specific/compare.mlir
@@ -0,0 +1,164 @@
+func @compare_tensor() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
+ %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
+ check.expect_eq_const(%output, dense<[0, 1, 0, 1]> : tensor<4xi8>) : tensor<4xi8>
+ return
+}
+
+func @compare_scalar() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<1> : tensor<i32>
+ %rhs = iree.unfoldable_constant dense<5> : tensor<i32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
+ check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+ return
+}
+
+func @compare_i8() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<1> : tensor<i8>
+ %rhs = iree.unfoldable_constant dense<5> : tensor<i8>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<i8>, tensor<i8>) -> tensor<i1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
+ check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+ return
+}
+
+func @compare_i16() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<1> : tensor<i16>
+ %rhs = iree.unfoldable_constant dense<5> : tensor<i16>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<i16>, tensor<i16>) -> tensor<i1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
+ check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+ return
+}
+
+func @compare_i32() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<1> : tensor<i32>
+ %rhs = iree.unfoldable_constant dense<5> : tensor<i32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
+ check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+ return
+}
+
+func @compare_i64() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<1> : tensor<i64>
+ %rhs = iree.unfoldable_constant dense<5> : tensor<i64>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<i64>, tensor<i64>) -> tensor<i1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
+ check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+ return
+}
+
+func @compare_f32() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<1.0> : tensor<f32>
+ %rhs = iree.unfoldable_constant dense<5.0> : tensor<f32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
+ check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+ return
+}
+
+func @compare_f64() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<1.0> : tensor<f64>
+ %rhs = iree.unfoldable_constant dense<5.0> : tensor<f64>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<f64>, tensor<f64>) -> tensor<i1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<i8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<i8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<i1>, tensor<i8>, tensor<i8>) -> tensor<i8>
+ check.expect_eq_const(%output, dense<0> : tensor<i8>) : tensor<i8>
+ return
+}
+
+func @compare_tensor_odd_length() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<[1, 2, 7]> : tensor<3xi32>
+ %rhs = iree.unfoldable_constant dense<[5, 2, 3]> : tensor<3xi32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<3xi8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<3xi8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<3xi1>, tensor<3xi8>, tensor<3xi8>) -> tensor<3xi8>
+ check.expect_eq_const(%output, dense<[0, 1, 0]> : tensor<3xi8>) : tensor<3xi8>
+ return
+}
+
+func @compare_eq() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
+ %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "EQ"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
+ check.expect_eq_const(%output, dense<[0, 1, 0, 1]> : tensor<4xi8>) : tensor<4xi8>
+ return
+}
+
+func @compare_ne() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
+ %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "NE"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
+ check.expect_eq_const(%output, dense<[1, 0, 1, 0]> : tensor<4xi8>) : tensor<4xi8>
+ return
+}
+
+func @compare_lt() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
+ %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "LT"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
+ check.expect_eq_const(%output, dense<[1, 0, 0, 0]> : tensor<4xi8>) : tensor<4xi8>
+ return
+}
+
+func @compare_le() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
+ %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "LE"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
+ check.expect_eq_const(%output, dense<[1, 1, 0, 1]> : tensor<4xi8>) : tensor<4xi8>
+ return
+}
+
+func @compare_gt() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
+ %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "GT"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
+ check.expect_eq_const(%output, dense<[0, 0, 1, 0]> : tensor<4xi8>) : tensor<4xi8>
+ return
+}
+
+func @compare_ge() attributes { iree.module.export } {
+ %lhs = iree.unfoldable_constant dense<[1, 2, 7, 4]> : tensor<4xi32>
+ %rhs = iree.unfoldable_constant dense<[5, 2, 3, 4]> : tensor<4xi32>
+ %result = "mhlo.compare"(%lhs, %rhs) {comparison_direction = "GE"} : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+ %c0 = iree.unfoldable_constant dense<0> : tensor<4xi8>
+ %c1 = iree.unfoldable_constant dense<1> : tensor<4xi8>
+ %output = "mhlo.select"(%result, %c1, %c0) : (tensor<4xi1>, tensor<4xi8>, tensor<4xi8>) -> tensor<4xi8>
+ check.expect_eq_const(%output, dense<[0, 1, 1, 1]> : tensor<4xi8>) : tensor<4xi8>
+ return
+}
diff --git a/iree/test/e2e/vulkan_specific/log_plus_one.mlir b/iree/test/e2e/vulkan_specific/log_plus_one.mlir
new file mode 100644
index 0000000..3bba4a9
--- /dev/null
+++ b/iree/test/e2e/vulkan_specific/log_plus_one.mlir
@@ -0,0 +1,6 @@
+func @log_plus_one() attributes { iree.module.export } {
+ %input = iree.unfoldable_constant dense<[0.0, 0.5, 1.0, 5.0]> : tensor<4xf32>
+ %result = "mhlo.log_plus_one"(%input) : (tensor<4xf32>) -> tensor<4xf32>
+ check.expect_almost_eq_const(%result, dense<[0.0, 0.4054651, 0.6931472, 1.7917595]> : tensor<4xf32>) : tensor<4xf32>
+ return
+}
diff --git a/scripts/get_e2e_artifacts.py b/scripts/get_e2e_artifacts.py
index 5cb6f8e..7ed2691 100755
--- a/scripts/get_e2e_artifacts.py
+++ b/scripts/get_e2e_artifacts.py
@@ -15,8 +15,12 @@
# limitations under the License.
"""Runs all E2E TensorFlow tests and extracts their benchmarking artifacts.
-Example usage:
- python3 get_e2e_artifacts.py
+Example usages:
+ # Run all test suites and collect their artifacts:
+ python3 ./scripts/get_e2e_artifacts.py
+
+ # Run the e2e_tests test suite and collect its artifacts:
+ python3 ./scripts/get_e2e_artifacts.py --test_suites=e2e_tests
"""
import fileinput
@@ -39,8 +43,10 @@
'//integrations/tensorflow/e2e:mobile_bert_squad_tests',
'keras_tests':
'//integrations/tensorflow/e2e/keras:keras_tests',
- 'vision_external_tests':
- '//integrations/tensorflow/e2e/keras:vision_external_tests',
+ 'imagenet_external_tests':
+ '//integrations/tensorflow/e2e/keras:imagenet_external_tests',
+ 'slim_vision_tests':
+ '//integrations/tensorflow/e2e/slim_vision_models:slim_vision_tests',
}
SUITES_HELP = [f'`{name}`' for name in SUITE_NAME_TO_TARGET]
SUITES_HELP = f'{", ".join(SUITES_HELP[:-1])} and {SUITES_HELP[-1]}'
@@ -118,6 +124,11 @@
paths_to_tests: Dict[str, str]):
"""Unzips all of the benchmarking artifacts for a given test and backend."""
outputs = os.path.join(test_path, 'test.outputs', 'outputs.zip')
+ if FLAGS.dry_run and not os.path.exists(outputs):
+ # The artifacts may or may not be present on disk during a dry run. If they
+ # are then we want to collision check them, but if they aren't that's fine.
+ return
+
archive = zipfile.ZipFile(outputs)
# Filter out directory names.
filenames = [name for name in archive.namelist() if name[-1] != os.sep]
@@ -139,17 +150,21 @@
# Convert test suite shorthands to full test suite targets.
test_suites = [SUITE_NAME_TO_TARGET[suite] for suite in FLAGS.test_suites]
+ if FLAGS.run_test_suites:
+ # Use bazel test to execute all of the test suites in parallel.
+ command = [
+ 'bazel', 'test', *test_suites, '--color=yes',
+ '--test_arg=--get_saved_model'
+ ]
+ print(f'Running: `{" ".join(command)}`')
+ if not FLAGS.dry_run:
+ subprocess.check_call(command)
+ print()
+
written_paths = set()
paths_to_tests = dict()
for test_suite in test_suites:
- if FLAGS.run_test_suites and not FLAGS.dry_run:
- subprocess.check_call([
- 'bazel', 'test', test_suite, '--color=yes',
- '--test_arg=--get_saved_model'
- ])
- print()
-
# Extract all of the artifacts for this test suite.
test_paths, test_names = get_test_paths_and_names(test_suite)
for i, (test_path, test_name) in enumerate(zip(test_paths, test_names)):
