Merge main -> google

* 82d8765 Merge branch 'google' into main-to-google
* b7e2b09 Remove special casing for IreeFileCheck in iree_lit_test.cmake (iree-org#5209)
* b67af5a Clean up several CMake options configured for our CI builds. (iree-org#5197)
* ea8b82d Integrate MLIR-EmitC at iml130/mlir-emitc@68547d0 (iree-org#5205)
* 10d9ea1 Make CPU/GPU end-to-end testing in Linalg on tensors/buffers align. (iree-org#5203)
* ff69d02 Move prepration of benchmark files to python script. (iree-org#5201)
* 0c2f9e8 Bump Abseil-cpp submodule to LTS 20200923, Patch 3 (iree-org#5198)
* 6ac0023 Add runtime tracing support to passes under IREE's PassManager usage. (iree-org#5160)
* 2e41217 Bump Benchmark submodule to 39c8d58 (iree-org#5196)
* 663ab93 Adjust UnixLinkerTool flags used for Apple targets. (iree-org#5164)
* 963beb9 Plumb pooling ops through Linalg on tensors path. (iree-org#5136)
* a33f72f Plumb depthwise convolution in Linalg on tensors. (iree-org#5189)
* 2864ee4 Add restrictions for enabling generic op vectorization. (iree-org#5188)
* 4f218a8 Make sure convolution vectorization is off when the option flag is off (iree-org#5187)
* 44736b5 build_tools:cmake:riscv: Fix the system library path (iree-org#5175)
* 96e9c49 Switching from mhlo::ConvertOp to ZeroExtendI for i1->i8 conversion. (iree-org#5119)
* da4ad51 Fix and re-enable jax testing (iree-org#5183)
* 301bc15 Merge google -> main (iree-org#5182)
* 400719c Shuffling around/removing some HLO deps.
* 9ba1d57 Enable MobileBert on GPU for the linalg on tensors path. (iree-org#5180)
* 52f02af Resolve name conflicts while linking executables in LLVM. (iree-org#5179)
* 4a732d3 Removing the post-partitioning conversion pass as it is unused. We've built up..

PiperOrigin-RevId: 364658113

CMakeLists.txt

@@ -38,7 +38,6 @@ option(IREE_BUILD_COMPILER "Builds the IREE compiler." ON)
 option(IREE_BUILD_TESTS "Builds IREE unit tests." ON)
 option(IREE_BUILD_DOCS "Builds IREE docs." OFF)
 option(IREE_BUILD_SAMPLES "Builds IREE sample projects." ON)
-option(IREE_BUILD_DEBUGGER "Builds the IREE debugger app." OFF)
 
 option(IREE_BUILD_TENSORFLOW_COMPILER "Builds TensorFlow compiler frontend." OFF)
 option(IREE_BUILD_TFLITE_COMPILER "Builds the TFLite compiler frontend." OFF)

SUBMODULE_VERSIONS.txt

@@ -1,12 +1,12 @@
-6ec136281086b71da32b5fb068bd6e46b78a5c79 third_party/abseil-cpp
-daff5fead3fbe22c6fc58310ca3f49caf117f185 third_party/benchmark
+6f9d96a1f41439ac172ee2ef7ccd8edf0e5d068c third_party/abseil-cpp
+39c8d58a76ce7c3398786ee96a452668e8392af7 third_party/benchmark
 5916273f79a21551890fd3d56fc5375a78d1598d third_party/cpuinfo
 4fb0ff7069bd88ee85902f4d0bb62794e5f6d021 third_party/flatcc
 b1fbd33c06cdb0024c67733c6fdec2009d17b384 third_party/googletest
 88b845dee001723c4a0db1fe5477de735b6d3bb0 third_party/liburing
 013b829185fee6d8eaa515a7e36ec468a2a02600 third_party/llvm-bazel
 b24436ac96bdf3f2c545fc85dc8af239d618c9c4 third_party/llvm-project
-3483f1653fc7cb3bfb3a4d1b463f3a651ecaa676 third_party/mlir-emitc
+68547d08daca039467df49c7cc50c3a0061787f3 third_party/mlir-emitc
 98debb127d3a14e0239a3432461e3876d293b409 third_party/mlir-hlo
 2b2bd45bbf9be04fd22ece5cc1f54679202e9257 third_party/pffft
 d8c7ee00a687ac369e62e2032514a93a9b413502 third_party/pybind11

build_tools/buildkite/cmake/android/arm64-v8a/benchmark.yml

@@ -16,13 +16,8 @@ steps:
   - label: "build"
     commands:
       - "docker run --user=$(id -u):$(id -g) --volume=\\$PWD:\\$IREE_DOCKER_WORKDIR --workdir=\\$IREE_DOCKER_WORKDIR --rm gcr.io/iree-oss/cmake-android@sha256:15d3266ae4865f7642a4ef4d76e5181f0dc3482a7cfba9021b6b55be524208ec build_tools/kokoro/gcp_ubuntu/cmake/android/build.sh arm64-v8a"
-      - "gsutil cp gs://iree-model-artifacts/iree-mobile-bert-artifacts-6fe4616e0ab9958eb18f368960a31276f1362029.tar.gz ."
-      - "gsutil cp gs://iree-model-artifacts/mobilenet-v2.tar.gz ."
-      - "tar -xzf iree-mobile-bert-artifacts-6fe4616e0ab9958eb18f368960a31276f1362029.tar.gz"
-      - "tar -xzf mobilenet-v2.tar.gz"
+      - "python3 build_tools/mako/prepare_benchmark_files.py"
       - "docker run --user=$(id -u):$(id -g) --volume=\\$PWD:\\$IREE_DOCKER_WORKDIR --workdir=\\$IREE_DOCKER_WORKDIR --rm gcr.io/iree-oss/cmake-android@sha256:15d3266ae4865f7642a4ef4d76e5181f0dc3482a7cfba9021b6b55be524208ec python3 build_tools/mako/compile_android_modules.py"
-      - "cp tmp/iree/modules/MobileBertSquad/iree_vmla/traces/serving_default/flagfile mobile-bert_flagfile"
-      - "cp mobilenet-v2/flagfile mobilenet-v2_flagfile"
       - "tar -czvf model-Pixel4-artifacts.tgz build-android/iree/tools/iree-benchmark-module *Pixel4*.vmfb"
       - "tar -czvf model-S20-artifacts.tgz build-android/iree/tools/iree-benchmark-module *S20*.vmfb"
       - "tar -czvf flagfiles.tgz *_flagfile"

build_tools/cmake/build_docs.sh

@@ -42,7 +42,6 @@ cd ${BUILD_DIR}
                 -DIREE_BUILD_TESTS=ON \
                 -DIREE_BUILD_SAMPLES=OFF \
                 -DIREE_BUILD_DOCS=ON \
-                -DIREE_BUILD_DEBUGGER=OFF \
                 -DIREE_BUILD_PYTHON_BINDINGS=OFF \
                 -G Ninja
 # Generate docs and also build iree-opt for generating more docs :)

build_tools/cmake/iree_lit_test.cmake

@@ -62,12 +62,7 @@ function(iree_lit_test)
   set(_DATA_DEP_PATHS)
   foreach(_DATA_DEP ${_RULE_DATA})
     string(REPLACE "::" "_" _DATA_DEP_NAME ${_DATA_DEP})
-    # TODO(*): iree_sh_binary so we can avoid this.
-    if("${_DATA_DEP_NAME}" STREQUAL "iree_tools_IreeFileCheck")
-      list(APPEND _DATA_DEP_PATHS "${CMAKE_SOURCE_DIR}/iree/tools/IreeFileCheck.sh")
-    else()
-      list(APPEND _DATA_DEP_PATHS $<TARGET_FILE:${_DATA_DEP_NAME}>)
-    endif()
+    list(APPEND _DATA_DEP_PATHS $<TARGET_FILE:${_DATA_DEP_NAME}>)
   endforeach(_DATA_DEP)
 
   iree_package_ns(_PACKAGE_NS)

build_tools/cmake/rebuild.sh

@@ -35,11 +35,19 @@ else
   mkdir build
 fi
 cd build
-"$CMAKE_BIN" -G Ninja -DCMAKE_BUILD_TYPE=RelWithDebInfo \
-                      -DIREE_BUILD_COMPILER=ON \
-                      -DIREE_BUILD_TESTS=ON \
-                      -DIREE_BUILD_SAMPLES=OFF \
-                      -DIREE_BUILD_DOCS=ON \
-                      -DIREE_BUILD_DEBUGGER=OFF \
-                      -DIREE_BUILD_PYTHON_BINDINGS=ON ..
+
+CMAKE_ARGS=(
+  "-G" "Ninja"
+  "-DCMAKE_BUILD_TYPE=RelWithDebInfo"
+
+  # Enable docs build on the CI. The additional builds are pretty fast and
+  # give us early warnings for some types of website publication errors.
+  "-DIREE_BUILD_DOCS=ON"
+
+  # Enable building the python bindings on CI. Most heavy targets are gated on
+  # IREE_ENABLE_TENSORFLOW, so what's left here should be fast.
+  "-DIREE_BUILD_PYTHON_BINDINGS=ON"
+)
+
+"$CMAKE_BIN" "${CMAKE_ARGS[@]?}" ..
 "$CMAKE_BIN" --build .

build_tools/cmake/riscv.toolchain.cmake

@@ -36,6 +36,7 @@ set(CMAKE_CXX_COMPILER "${RISCV_TOOLCHAIN_ROOT}/bin/clang++")
 set(CMAKE_AR "${RISCV_TOOLCHAIN_ROOT}/bin/llvm-ar")
 set(CMAKE_RANLIB "${RISCV_TOOLCHAIN_ROOT}/bin/llvm-ranlib")
 set(CMAKE_STRIP "${RISCV_TOOLCHAIN_ROOT}/bin/llvm-strip")
+set(CMAKE_SYSTEM_LIBRARY_PATH "${RISCV_TOOLCHAIN_ROOT}/sysroot/usr/lib")
 
 set(RISCV_COMPILER_FLAGS)
 set(RISCV_COMPILER_FLAGS_CXX)

build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-swiftshader/build.sh

@@ -58,7 +58,6 @@ echo "Configuring CMake"
    -DIREE_BUILD_COMPILER=ON \
    -DIREE_BUILD_TESTS=ON \
    -DIREE_BUILD_SAMPLES=OFF \
-   -DIREE_BUILD_DEBUGGER=OFF \
    -DIREE_BUILD_XLA_COMPILER=ON \
    -DIREE_BUILD_TFLITE_COMPILER=ON \
    -DIREE_BUILD_TENSORFLOW_COMPILER=ON .
@@ -70,4 +69,6 @@ ninja
 export CTEST_PARALLEL_LEVEL=${CTEST_PARALLEL_LEVEL:-$(nproc)}
 
 echo "Testing with CTest"
-ctest --output-on-failure -L 'integrations/tensorflow' --label-exclude "^nokokoro$"
+ctest --output-on-failure \
+   --tests-regex "^integrations/tensorflow/|^bindings/python/" \
+   --label-exclude "^nokokoro$"

build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-turing/build.sh

@@ -61,7 +61,6 @@ echo "Configuring CMake"
    -DIREE_BUILD_COMPILER=ON \
    -DIREE_BUILD_TESTS=ON \
    -DIREE_BUILD_SAMPLES=OFF \
-   -DIREE_BUILD_DEBUGGER=OFF \
    -DIREE_BUILD_XLA_COMPILER=ON \
    -DIREE_BUILD_TFLITE_COMPILER=ON \
    -DIREE_BUILD_TENSORFLOW_COMPILER=ON .
@@ -78,6 +77,6 @@ export CTEST_PARALLEL_LEVEL=${CTEST_PARALLEL_LEVEL:-1}
 # as well.
 echo "Testing with CTest"
 ctest --output-on-failure \
-   --tests-regex "^integrations/tensorflow/" \
+   --tests-regex "^integrations/tensorflow/|^bindings/python/" \
    --label-regex "^driver=vulkan$|^driver=cuda$" \
    --label-exclude "^nokokoro$"

build_tools/mako/benchmark_modules_on_android.py

@@ -88,7 +88,7 @@ def main(args) -> None:
       for target in phone.targets:
         module_name = configuration.get_module_name(model_benchmark.name,
                                                     phone.name, target.mako_tag)
-        flagfile_name = f"{model_benchmark.name}_flagfile"
+        flagfile_name = configuration.get_flagfile_name(model_benchmark.name)
         mako_log.append(benchmark(module_name, flagfile_name, target))
       mako_log.append(
           get_mako_metadata(args.git_hash, timestamp, phone.benchmark_key))

build_tools/mako/configuration.py

@@ -76,15 +76,21 @@ class ModelBenchmarkInfo:
 
   Attributes:
     name: The name of the model.
+    model_artifacts_name: The filename of model artifacts which locates in
+      Google bucket iree-model-artifacts.
     model_path: A path to MLIR input file. This can be a relative path.
+    flagfile_path: A path to flagfile. This can be a relative path.
     phones: A list of PhoneBenchmarkInfo that the benchmarking targets on.
   """
 
-  def __init__(self, name, model_path, phones):
+  def __init__(self, name, model_artifacts_name, model_path, flagfile_path,
+               phones):
     if "_" in name:
       raise ValueError("The target name contains invalid char '_'")
     self.name = name
+    self.model_artifacts_name = model_artifacts_name
     self.model_path = model_path
+    self.flagfile_path = flagfile_path
     self.phones = phones
 
 
@@ -144,7 +150,11 @@ def get_s20_default_target_list(batch_config=None):
 MODEL_BENCHMARKS = [
     ModelBenchmarkInfo(
         name="mobile-bert",
+        model_artifacts_name=
+        "iree-mobile-bert-artifacts-6fe4616e0ab9958eb18f368960a31276f1362029.tar.gz",
         model_path="tmp/iree/modules/MobileBertSquad/iree_input.mlir",
+        flagfile_path=
+        "tmp/iree/modules/MobileBertSquad/iree_vmla/traces/serving_default/flagfile",
         phones=[
             PhoneBenchmarkInfo(name="Pixel4",
                                benchmark_key="5538704950034432",
@@ -159,7 +169,9 @@ def get_s20_default_target_list(batch_config=None):
         ]),
     ModelBenchmarkInfo(
         name="mobilenet-v2",
+        model_artifacts_name="mobilenet-v2.tar.gz",
         model_path="mobilenet-v2/iree_input.mlir",
+        flagfile_path="mobilenet-v2/flagfile",
         phones=[
             PhoneBenchmarkInfo(name="Pixel4",
                                benchmark_key="6338759231537152",
@@ -175,5 +187,9 @@ def get_s20_default_target_list(batch_config=None):
 ]
 
 
+def get_flagfile_name(model_name):
+  return f"{model_name}_flagfile"
+
+
 def get_module_name(model_name, phone_name, mako_tag):
   return f"{model_name}_{phone_name}_{mako_tag}.vmfb"

build_tools/mako/prepare_benchmark_files.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+# 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.
+
+# Prepares files for
+#   * compile_android_modules.py
+#   * benchmark_modules_on_android.py
+#
+# The script assumes model artifacts are in google bucket and they are zipped in
+# tar.gz format.
+
+import subprocess
+
+import configuration
+
+
+def main() -> None:
+  for model_benchmark in configuration.MODEL_BENCHMARKS:
+    print(f"Preparing benchmark files for {model_benchmark.name}")
+    subprocess.run(args=[
+        "gsutil", "cp",
+        f"gs://iree-model-artifacts/{model_benchmark.model_artifacts_name}", "."
+    ],
+                   check=True)
+    subprocess.run(args=["tar", "-xvf", model_benchmark.model_artifacts_name],
+                   check=True)
+    subprocess.run(args=[
+        "cp", model_benchmark.flagfile_path,
+        configuration.get_flagfile_name(model_benchmark.name)
+    ],
+                   check=True)
+
+
+if __name__ == "__main__":
+  main()

docs/get_started/cmake_options_and_variables.md

@@ -48,10 +48,6 @@ Builds IREE documentation. Defaults to `OFF`.
 
 Builds IREE sample projects. Defaults to `ON`.
 
-#### `IREE_BUILD_DEBUGGER`:BOOL
-
-Builds the IREE debugger app. Defaults to `OFF`.
-
 #### `IREE_BUILD_PYTHON_BINDINGS`:BOOL
 
 Builds the IREE python bindings. Defaults to `OFF`.

integrations/tensorflow/iree_tf_compiler/BUILD

@@ -93,6 +93,7 @@ cc_binary(
         "//iree_tf_compiler/TF",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",
         "@org_tensorflow//tensorflow/compiler/mlir/xla:hlo_to_mlir_hlo",
         "@org_tensorflow//tensorflow/compiler/xla/service:hlo_parser",

integrations/tensorflow/iree_tf_compiler/TF/Passes.cpp

@@ -127,6 +127,10 @@ void buildTFImportPassPipeline(OpPassManager &pm) {
   pm.nest<ModuleOp>().addPass(createStripFunctionMetadataPass());
   pm.addPass(createVerifyFullyConvertedPass());
 
+  buildMHLOImportPassPipeline(pm);
+}
+
+void buildMHLOImportPassPipeline(OpPassManager &pm) {
   //----------------------------------------------------------------------------
   // Convert control flow and flatten tuples (like tuple<tensor<...>, ...>)
   //----------------------------------------------------------------------------
@@ -148,6 +152,15 @@ void buildTFImportPassPipeline(OpPassManager &pm) {
   pm.addPass(iree_compiler::Shape::createConvertHLOToShapePass());
 }
 
+void registerMHLOImportPassPipeline() {
+  mlir::PassPipelineRegistration<> pipeline(
+      "iree-mhlo-import-pipeline",
+      "Run IREE-specific passes for importing MHLO code into IREE",
+      [](OpPassManager &passManager) {
+        buildMHLOImportPassPipeline(passManager);
+      });
+}
+
 void registerTFImportPassPipeline() {
   mlir::PassPipelineRegistration<> pipeline(
       "iree-tf-import-pipeline",

integrations/tensorflow/iree_tf_compiler/TF/Passes.h

@@ -33,8 +33,12 @@ namespace TF {
 // passes in the right order.
 void buildTFImportPassPipeline(OpPassManager &pm);
 
+void buildMHLOImportPassPipeline(OpPassManager &pm);
+
 void registerTFImportPassPipeline();
 
+void registerMHLOImportPassPipeline();
+
 //===----------------------------------------------------------------------===//
 // IREE-specific Passes For TensorFlow Import
 //===----------------------------------------------------------------------===//
@@ -75,6 +79,7 @@ void registerAllDialects(mlir::DialectRegistry &registry);
 
 inline void registerAllPasses() {
   registerTFImportPassPipeline();
+  registerMHLOImportPassPipeline();
 
   createConvertToMHLOPass();
   createFlattenTuplesInCFGPass();

integrations/tensorflow/iree_tf_compiler/iree-import-xla-main.cpp

@@ -18,6 +18,7 @@
 #include <fstream>
 #include <iostream>
 
+#include "iree_tf_compiler/TF/Passes.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/InitLLVM.h"
@@ -28,6 +29,7 @@
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/PassManager.h"
 #include "mlir/Support/FileUtilities.h"
 #include "tensorflow/compiler/mlir/xla/hlo_to_mlir_hlo.h"
 #include "tensorflow/compiler/xla/service/hlo.pb.h"
@@ -218,6 +220,17 @@ int main(int argc, char **argv) {
     if (failed(saveToFile(saveTempIreeImport))) return 10;
   }
 
+  // Run passes.
+  PassManager pm(&context, PassManager::Nesting::Implicit);
+  applyPassManagerCLOptions(pm);
+
+  iree_integrations::TF::buildMHLOImportPassPipeline(pm);
+  if (failed(pm.run(*module))) {
+    llvm::errs()
+        << "Running iree-xla-import pass pipeline failed (see diagnostics)\n";
+    return 2;
+  }
+
   // Save output.
   if (failed(saveToFile(outputFilename))) return 3;
   return 0;

iree/compiler/Conversion/Common/LinalgBufferizePass.cpp

@@ -232,6 +232,22 @@ static LogicalResult convertConstantOp(OpBuilder &b, ConstantOp constantOp,
   return success();
 }
 
+/// Converts a linalg.init_tensor op to memref.alloc op. This provides a shaped
+/// operand for pooling ops. The op will be deleted after going to loops.
+static LogicalResult convertInitTensorOp(
+    OpBuilder &b, WorkgroupMemoryAllocationFn allocationFn,
+    linalg::InitTensorOp initTensorOp, BlockAndValueMapping &bvm) {
+  if (bvm.contains(initTensorOp.getResult())) return success();
+  RankedTensorType tensorType = initTensorOp.getType();
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPointAfter(initTensorOp);
+  Value alloc = allocationFn(b, initTensorOp.getLoc(), tensorType.getShape(),
+                             tensorType.getElementType(),
+                             llvm::to_vector<4>(initTensorOp.sizes()));
+  bvm.map(initTensorOp.getResult(), alloc);
+  return success();
+}
+
 /// Avoids creating an allocation if the result tensor can just be aliased to
 /// use the same buffer (`inputBuffer`) that `srcTensor` is mapped to. This can
 /// be done if `srcTensor` has a single use, which is the operation which is
@@ -668,6 +684,9 @@ void LinalgBufferizePass::runOnFunction() {
         .Case<linalg::TensorReshapeOp>([&](linalg::TensorReshapeOp reshapeOp) {
           return convertTensorReshapeOp(b, allocationFn, reshapeOp, bvm);
         })
+        .Case<linalg::InitTensorOp>([&](linalg::InitTensorOp initTensorOp) {
+          return convertInitTensorOp(b, allocationFn, initTensorOp, bvm);
+        })
         .Case<tensor::ExtractOp>([&](tensor::ExtractOp extractOp) {
           return convertTensorExtractOp(b, extractOp, bvm);
         })

iree/compiler/Conversion/Common/Transforms.cpp

@@ -317,10 +317,9 @@ LogicalResult defineWorkgroupCountRegion(
     OpBuilder &builder, FuncOp funcOp,
     WorkgroupCountRegionBuilder regionBuilder) {
   IREE::HAL::ExecutableEntryPointOp entryPointOp = getEntryPoint(funcOp);
-  if (!entryPointOp)
+  if (!entryPointOp) {
     return funcOp.emitOpError("unable to find corresponding entry point op");
-  if (entryPointOp.getBody())
-    return entryPointOp.emitOpError("cannot override workgroup_count_region");
+  }
   Location loc = entryPointOp.getLoc();
 
   OpBuilder::InsertionGuard guard(builder);