Merge pull request iree-org#5078 from ThomasRaoux:main-to-google

PiperOrigin-RevId: 362543741
not-jenni · Mar 12, 2021 · c9e5da3 · c9e5da3
2 parents 284e68c + e1136e3
commit c9e5da3
Show file tree

Hide file tree

Showing 303 changed files with 6,463 additions and 3,176 deletions.
diff --git a/build_tools/bazel/deep_copy.bzl b/build_tools/bazel/deep_copy.bzl
@@ -43,6 +43,7 @@ def _deep_copy_recursion_depth_1(x):
 
 def deep_copy(x):
     """Returns a copy of the argument, making a deep copy if it is a container.
+
     Args:
       x: (object) value to copy. If it is a container with nested containers as
          elements, the maximum nesting depth is restricted to three (e.g.,

diff --git a/build_tools/bazel/enforce_glob.bzl b/build_tools/bazel/enforce_glob.bzl
@@ -0,0 +1,59 @@
+# Copyright 2021 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.
+
+"""A utility to enforce that a list matches a glob expression.
+
+We use this primarily to enable the error-checking capabilities of globs of test
+files in IREE while still allowing our Bazel to CMake conversion to not create
+CMake globs (which are discouraged for collecting source files, see
+https://cmake.org/cmake/help/latest/command/file.html#glob) and not be dependent
+on any information outside of the BUILD file.
+"""
+
+def enforce_glob(files, **kwargs):
+    """A utility to enforce that a list matches a glob expression.
+
+    Note that the comparison is done in an order-independent fashion.
+
+    Args:
+        files: a list that is expected to contain the same files as the
+            specified glob expression.
+        **kwargs: keyword arguments forwarded to the glob.
+
+    Returns:
+        files. The input argument unchanged
+    """
+    glob_result = native.glob(**kwargs)
+
+    # glob returns a sorted list.
+    if sorted(files) != glob_result:
+        glob_result_dict = {k: None for k in glob_result}
+        result_dict = {k: None for k in files}
+        missing = [k for k in glob_result if k not in files]
+        extra = [k for k in files if k not in glob_result]
+        expected_formatted = "\n".join(['"{}",'.format(file) for file in glob_result])
+        fail(("Error in enforce_glob." +
+              "\nExpected {}." +
+              "\nGot {}." +
+              "\nMissing {}." +
+              "\nExtra {}" +
+              "\nPaste this into the first enforce_glob argument:" +
+              "\n{}").format(
+            glob_result,
+            files,
+            missing,
+            extra,
+            expected_formatted,
+        ))
+    return files
diff --git a/build_tools/bazel_to_cmake/bazel_to_cmake_converter.py b/build_tools/bazel_to_cmake/bazel_to_cmake_converter.py
@@ -251,6 +251,9 @@ def filegroup(self, name, **kwargs):
   def sh_binary(self, name, **kwargs):
     self._convert_unimplemented_function("sh_binary", name)
 
+  def enforce_glob(self, files, **kwargs):
+    return files
+
   def glob(self, include, exclude=None, exclude_directories=1):
     if exclude_directories != 1:
       self._convert_unimplemented_function("glob", "with exclude_directories")

diff --git a/build_tools/docker/cmake-bazel-frontends-nvidia/Dockerfile b/build_tools/docker/cmake-bazel-frontends-nvidia/Dockerfile
@@ -26,4 +26,6 @@
 FROM gcr.io/iree-oss/cmake-bazel-frontends-vulkan@sha256:3656d9c3a08770f8371ad3a6f777b979c954656dae1cdbdfadfccdd7ab713f87 AS final
 
 RUN apt-get update \
-  && DEBIAN_FRONTEND=noninteractive apt-get install -y libnvidia-gl-460=460.32.03-0ubuntu0.18.04.1
+  && DEBIAN_FRONTEND=noninteractive apt-get install -y \
+       libnvidia-gl-460=460.39-0ubuntu0.18.04.1 \
+       libnvidia-compute-460=460.39-0ubuntu0.18.04.1
diff --git a/build_tools/docker/cmake-python-nvidia/Dockerfile b/build_tools/docker/cmake-python-nvidia/Dockerfile
@@ -28,4 +28,6 @@
 FROM gcr.io/iree-oss/cmake-python-vulkan@sha256:6722f69c6300749f6bd4b141fc653244990381a6b0111f9c361061adcd65c07c AS final
 
 RUN apt-get update \
-  && DEBIAN_FRONTEND=noninteractive apt-get install -y libnvidia-gl-460=460.32.03-0ubuntu0.18.04.1
+  && DEBIAN_FRONTEND=noninteractive apt-get install -y \
+       libnvidia-gl-460=460.39-0ubuntu0.18.04.1 \
+       libnvidia-compute-460=460.39-0ubuntu0.18.04.1
diff --git a/build_tools/docker/prod_digests.txt b/build_tools/docker/prod_digests.txt
@@ -9,8 +9,8 @@ gcr.io/iree-oss/vulkan@sha256:5812ee64806a7f3df0739ccf0930c27cabce346901488eceb1
 gcr.io/iree-oss/rbe-toolchain@sha256:d69c260b98a97ad430d34c4591fb2399e00888750f5d47ede00c1e6f3e774e5a
 gcr.io/iree-oss/cmake-python-vulkan@sha256:6722f69c6300749f6bd4b141fc653244990381a6b0111f9c361061adcd65c07c
 gcr.io/iree-oss/cmake-python-swiftshader@sha256:0be2b0c735a038365e7cad31f6b440805dd4e231e166a114ef22914a5469cbc8
-gcr.io/iree-oss/cmake-python-nvidia@sha256:2d823e6fc528d0da5296ae35c0aad8008306381e3ac923679c4052adedd029d7
+gcr.io/iree-oss/cmake-python-nvidia@sha256:0c931cac303791af85c5a717b418997cac9f3319717f59f7a70ac777edfa7b33
 gcr.io/iree-oss/cmake-bazel-frontends@sha256:2a9c65cea8b061696c8217b78e9718ceb804b9245bb3a661ff8a56f832aeeb0a
 gcr.io/iree-oss/cmake-bazel-frontends-vulkan@sha256:3656d9c3a08770f8371ad3a6f777b979c954656dae1cdbdfadfccdd7ab713f87
-gcr.io/iree-oss/cmake-bazel-frontends-nvidia@sha256:e71d77c2e8b886e99ffce2f50c73c5ae5e316bde23b54f0e6090bb04ce2d1bdf
+gcr.io/iree-oss/cmake-bazel-frontends-nvidia@sha256:bc1502af0679d301feb27477c2476006282b5f956cecb650b2716c7a2876e722
 gcr.io/iree-oss/cmake-bazel-frontends-swiftshader@sha256:1466d3658f872a0675b3ac605c2e577def880c7334738d677d5ef2f1c1291646
diff --git a/build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-turing/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-turing/build_kokoro.sh
@@ -38,7 +38,7 @@ docker_setup
 
 docker run "${DOCKER_RUN_ARGS[@]?}" \
   --gpus all \
-  gcr.io/iree-oss/cmake-bazel-frontends-nvidia@sha256:e71d77c2e8b886e99ffce2f50c73c5ae5e316bde23b54f0e6090bb04ce2d1bdf \
+  gcr.io/iree-oss/cmake-bazel-frontends-nvidia@sha256:bc1502af0679d301feb27477c2476006282b5f956cecb650b2716c7a2876e722 \
   build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-turing/build.sh
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build_kokoro.sh
@@ -38,7 +38,7 @@ docker_setup
 
 docker run "${DOCKER_RUN_ARGS[@]?}" \
   --gpus all \
-  gcr.io/iree-oss/cmake-python-nvidia@sha256:2d823e6fc528d0da5296ae35c0aad8008306381e3ac923679c4052adedd029d7 \
+  gcr.io/iree-oss/cmake-python-nvidia@sha256:0c931cac303791af85c5a717b418997cac9f3319717f59f7a70ac777edfa7b33 \
   build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build.sh
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/integrations/tensorflow/iree_tf_compiler/TF/test/BUILD b/integrations/tensorflow/iree_tf_compiler/TF/test/BUILD
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("@iree//iree:lit_test.bzl", "iree_lit_test_suite")
+load("@iree//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -22,7 +23,18 @@ package(
 
 iree_lit_test_suite(
     name = "lit",
-    srcs = glob(["*.mlir"]),
+    srcs = enforce_glob(
+        [
+            "convert_to_mhlo.mlir",
+            "lower_global_tensors.mlir",
+            "lower_global_tensors_complex.mlir",
+            "lower_global_tensors_invalid.mlir",
+            "propagate_resource_casts.mlir",
+            "strip_metadata.mlir",
+            "verify_fully_converted.mlir",
+        ],
+        include = ["*.mlir"],
+    ),
     data = [
         "//iree_tf_compiler:iree-tf-opt",
         "@iree//iree/tools:IreeFileCheck",

diff --git a/integrations/tensorflow/iree_tf_compiler/TFL/test/BUILD b/integrations/tensorflow/iree_tf_compiler/TFL/test/BUILD
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("@iree//iree:lit_test.bzl", "iree_lit_test_suite")
+load("@iree//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -22,7 +23,14 @@ package(
 
 iree_lit_test_suite(
     name = "lit",
-    srcs = glob(["*.mlir"]),
+    srcs = enforce_glob(
+        [
+            "convert_metadata.mlir",
+            "strip_metadata.mlir",
+            "verify_fully_converted.mlir",
+        ],
+        include = ["*.mlir"],
+    ),
     data = [
         "//iree_tf_compiler:iree-opt-tflite",
         "@iree//iree/tools:IreeFileCheck",

diff --git a/integrations/tensorflow/iree_tf_compiler/TFL/test/import/BUILD b/integrations/tensorflow/iree_tf_compiler/TFL/test/import/BUILD
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("@iree//iree:lit_test.bzl", "iree_lit_test_suite")
+load("@iree//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -22,7 +23,13 @@ package(
 
 iree_lit_test_suite(
     name = "lit",
-    srcs = glob(["*.mlir"]),
+    srcs = enforce_glob(
+        [
+            "add.mlir",
+            "multi_add.mlir",
+        ],
+        include = ["*.mlir"],
+    ),
     data = glob(["*.tflite"]) + [
         "//iree_tf_compiler:iree-import-tflite",
         "@iree//iree/tools:IreeFileCheck",

diff --git a/integrations/tensorflow/iree_tf_compiler/dialect/tf_strings/conversion/test/BUILD b/integrations/tensorflow/iree_tf_compiler/dialect/tf_strings/conversion/test/BUILD
@@ -15,6 +15,7 @@
 # Tests for lowering MLIR in various dialects to IREE interpreter bytecode.
 
 load("@iree//iree:lit_test.bzl", "iree_lit_test_suite")
+load("@iree//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -24,7 +25,13 @@ package(
 
 iree_lit_test_suite(
     name = "lit",
-    srcs = glob(["*.mlir"]),
+    srcs = enforce_glob(
+        [
+            "tf_strings_to_strings.mlir",
+            "tf_to_tf_strings.mlir",
+        ],
+        include = ["*.mlir"],
+    ),
     data = [
         "//iree_tf_compiler:iree-tf-opt",
         "@iree//iree/tools:IreeFileCheck",

diff --git a/integrations/tensorflow/iree_tf_compiler/dialect/tf_strings/ir/BUILD b/integrations/tensorflow/iree_tf_compiler/dialect/tf_strings/ir/BUILD
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("@iree//build_tools/bazel:tblgen.bzl", "gentbl")
+load("@iree//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -22,7 +23,13 @@ package(
 
 filegroup(
     name = "td_files",
-    srcs = glob(["*.td"]),
+    srcs = enforce_glob(
+        [
+            "base.td",
+            "ops.td",
+        ],
+        include = ["*.td"],
+    ),
 )
 
 gentbl(

diff --git a/integrations/tensorflow/iree_tf_compiler/dialect/tf_tensorlist/conversion/test/BUILD b/integrations/tensorflow/iree_tf_compiler/dialect/tf_tensorlist/conversion/test/BUILD
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("@iree//iree:lit_test.bzl", "iree_lit_test_suite")
+load("@iree//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -22,7 +23,13 @@ package(
 
 iree_lit_test_suite(
     name = "lit",
-    srcs = glob(["*.mlir"]),
+    srcs = enforce_glob(
+        [
+            "convert_tf_tensorlist_to_tensorlist.mlir",
+            "convert_tf_to_tf_tensorlist.mlir",
+        ],
+        include = ["*.mlir"],
+    ),
     data = [
         "//iree_tf_compiler:iree-tf-opt",
         "@iree//iree/tools:IreeFileCheck",

diff --git a/integrations/tensorflow/iree_tf_compiler/dialect/tf_tensorlist/ir/BUILD b/integrations/tensorflow/iree_tf_compiler/dialect/tf_tensorlist/ir/BUILD
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("@iree//build_tools/bazel:tblgen.bzl", "gentbl")
+load("@iree//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -24,7 +25,13 @@ exports_files(["tf_tensorlist_base.td"])
 
 filegroup(
     name = "td_files",
-    srcs = glob(["*.td"]),
+    srcs = enforce_glob(
+        [
+            "tf_tensorlist_base.td",
+            "tf_tensorlist_ops.td",
+        ],
+        include = ["*.td"],
+    ),
 )
 
 cc_library(

diff --git a/integrations/tensorflow/iree_tf_compiler/dialect/tf_tensorlist/ir/test/BUILD b/integrations/tensorflow/iree_tf_compiler/dialect/tf_tensorlist/ir/test/BUILD
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("@iree//iree:lit_test.bzl", "iree_lit_test_suite")
+load("@iree//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -22,7 +23,10 @@ package(
 
 iree_lit_test_suite(
     name = "lit",
-    srcs = glob(["*.mlir"]),
+    srcs = enforce_glob(
+        ["ops.mlir"],
+        include = ["*.mlir"],
+    ),
     data = [
         "//iree_tf_compiler:iree-tf-opt",
         "@iree//iree/tools:IreeFileCheck",

diff --git a/iree/compiler/Bindings/TFLite/Transforms/test/BUILD b/iree/compiler/Bindings/TFLite/Transforms/test/BUILD
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -22,7 +23,13 @@ package(
 
 iree_lit_test_suite(
     name = "lit",
-    srcs = glob(["*.mlir"]),
+    srcs = enforce_glob(
+        [
+            "materialize_shape_support.mlir",
+            "wrap_entry_points.mlir",
+        ],
+        include = ["*.mlir"],
+    ),
     data = [
         "//iree/tools:IreeFileCheck",
         "//iree/tools:iree-opt",

diff --git a/iree/compiler/Bindings/TFLite/Transforms/test/CMakeLists.txt b/iree/compiler/Bindings/TFLite/Transforms/test/CMakeLists.txt
@@ -10,12 +10,12 @@
 
 iree_add_all_subdirs()
 
-file(GLOB _GLOB_X_MLIR LIST_DIRECTORIES false RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} CONFIGURE_DEPENDS *.mlir)
 iree_lit_test_suite(
   NAME
     lit
   SRCS
-    "${_GLOB_X_MLIR}"
+    "materialize_shape_support.mlir"
+    "wrap_entry_points.mlir"
   DATA
     iree::tools::IreeFileCheck
     iree::tools::iree-opt

diff --git a/iree/compiler/Conversion/Common/LaunchConfig.cpp b/iree/compiler/Conversion/Common/LaunchConfig.cpp
@@ -28,6 +28,7 @@
 #include "llvm/Support/FormatVariadic.h"
 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Operation.h"
@@ -38,6 +39,7 @@ namespace iree_compiler {
 /// Name of the StrAttr that can be used to get the key to access the tile size
 /// information.
 static const char kLaunchInfoKey[] = "launch_info_key";
+static const char kRootOpKey[] = "is_root_op";
 
 static Optional<StringRef> getKey(Operation *op) {
   StringAttr attr = op->getAttrOfType<StringAttr>(kLaunchInfoKey);
@@ -82,8 +84,7 @@ ArrayRef<int64_t> LaunchConfig::getTileSizes(Operation *op,
 
 Operation *LaunchConfig::getRootOperation(ArrayRef<Operation *> ops) {
   for (auto op : ops) {
-    auto key = getKey(op);
-    if (key && key.getValue() == rootOperationKey) return op;
+    if (op->getAttrOfType<UnitAttr>(kRootOpKey)) return op;
   }
   return nullptr;
 }
@@ -114,8 +115,7 @@ void LaunchConfig::setNumSubgroups(ArrayRef<int64_t> vNumSubgroups) {
 }
 
 void LaunchConfig::setRootOperation(Operation *op) {
-  Optional<StringRef> key = getKey(op);
-  if (key) rootOperationKey = *key;
+  op->setAttr(kRootOpKey, UnitAttr::get(op->getContext()));
 }
 
 void LaunchConfig::setSameConfig(Operation *source, Operation *target) {

diff --git a/iree/compiler/Conversion/Common/LaunchConfig.h b/iree/compiler/Conversion/Common/LaunchConfig.h
@@ -127,10 +127,6 @@ class LaunchConfig {
   /// these attributes.
   llvm::StringMap<TileSizesListType> tileSizes;
 
-  /// Key used for tagging the root operation. The launch config does not track
-  /// the root operation itself, but rather the key used for the root operation.
-  StringRef rootOperationKey = "";
-
   /// Workgroup size to use.
   std::array<int64_t, 3> workgroupSize = {1, 1, 1};