chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/lite/debug/BUILD

@@ -0,0 +1,69 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load(
+    "//tensorflow/core/platform:build_config.bzl",
+    "tf_proto_library",
+)
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
+
+# OSS only: This target is header-only. Link `debug_options_proto_cc_impl` only to
+# `libtensorflow_framework.so` via `lib_internal_impl`. Do NOT link `debug_options_proto_cc_impl`
+# directly unless the target does not link `libtensorflow_framework.so`.
+tf_proto_library(
+    name = "debug_options_proto",
+    srcs = ["debug_options.proto"],
+    make_default_target_header_only = True,
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "debug",
+    srcs = ["debug.cc"],
+    hdrs = ["debug.h"],
+    visibility = ["//tensorflow/compiler/mlir/lite:__subpackages__"],
+    deps = [
+        ":debug_options_proto_cc",
+        "//tensorflow/compiler/mlir/lite/metrics:error_collector_inst",
+        "//tensorflow/core:portable_gif_internal",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
+        "@com_googlesource_code_re2//:re2",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+        "@tsl//tsl/platform:path",
+        "@tsl//tsl/platform:stringpiece",
+        "@xla//xla/tsl/lib/io:buffered_file",
+        "@xla//xla/tsl/platform:env",
+    ],
+)
+
+tf_cc_test(
+    name = "debug_test",
+    srcs = ["debug_test.cc"],
+    deps = [
+        ":debug",
+        ":debug_options_proto_cc",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/core:portable_gif_internal",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@tsl//tsl/platform:path",
+        "@xla//xla/tsl/lib/core:status_test_util",
+        "@xla//xla/tsl/platform:env",
+    ],
+)

tensorflow/compiler/mlir/lite/debug/debug.cc

@@ -0,0 +1,358 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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 "tensorflow/compiler/mlir/lite/debug/debug.h"
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassInstrumentation.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/FileUtilities.h"  // from @llvm-project
+#include "mlir/Transforms/ViewOpGraph.h"  // from @llvm-project
+#include "re2/re2.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"
+#include "tensorflow/compiler/mlir/lite/metrics/error_collector_inst.h"
+#include "xla/tsl/lib/io/buffered_file.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/file_system.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tsl/platform/path.h"
+#include "tsl/platform/stringpiece.h"
+
+// IWYU pragma: no_include "util/regexp/re2/re2.h"
+
+namespace tensorflow {
+namespace {
+
+// Simple raw_ostream that prints to a file.
+struct WritableFileRawStream : public llvm::raw_ostream {
+  explicit WritableFileRawStream(std::unique_ptr<tsl::WritableFile> file)
+      : file(std::move(file)) {
+    SetUnbuffered();
+  }
+  ~WritableFileRawStream() override = default;
+
+  uint64_t current_pos() const override {
+    int64_t position;
+    if (file->Tell(&position).ok()) {
+      return position;
+    } else {
+      // MLIR uses os.tell() to determine whether something was written by
+      // a subroutine or not, so it's important we have a working current_pos().
+      LOG(WARNING)
+          << "Couldn't query file position. Stream might be malformed.\n";
+      return -1;
+    }
+  }
+
+  void write_impl(const char* ptr, size_t size) override {
+    // Write the file if it is still valid. If the write fails, null out the
+    // file to avoid encountering another error.
+    if (file && !file->Append(absl::string_view(ptr, size)).ok()) {
+      file = nullptr;
+    }
+  }
+
+  // The file being written to.
+  std::unique_ptr<tsl::WritableFile> file;
+};
+
+// Reproducer stream that emits a reproducer to the given `llvm::raw_ostream`.
+class ReproducerStream : public mlir::ReproducerStream {
+ public:
+  ReproducerStream(std::string name, std::unique_ptr<llvm::raw_ostream> os)
+      : name_(std::move(name)), os_(std::move(os)) {}
+
+  llvm::StringRef description() override { return name_; }
+
+  llvm::raw_ostream& os() override { return *os_; }
+
+ private:
+  std::string name_;
+  std::unique_ptr<llvm::raw_ostream> os_;
+};
+
+// Returns a function that builds a reproducer stream, or nullptr if the MLIR
+// reproducer will not be enabled.
+mlir::ReproducerStreamFactory GetReproducerStreamFactory(
+    absl::string_view dump_dir) {
+  std::string path = tsl::io::JoinPath(dump_dir, "tfl_mlir_crash_repro.mlir");
+
+  return [path = std::move(path)](
+             std::string& error) -> std::unique_ptr<mlir::ReproducerStream> {
+    std::unique_ptr<tsl::WritableFile> file;
+    if (auto status = tsl::Env::Default()->NewWritableFile(path, &file);
+        !status.ok()) {
+      error = status.ToString();
+      absl::StrAppend(&error, "; failed to open '", path,
+                      "' for writing an MLIR reproducer");
+      return nullptr;
+    }
+    file = std::make_unique<tsl::BufferedWritableFile>(std::move(file));
+
+    return std::make_unique<ReproducerStream>(
+        path, std::make_unique<WritableFileRawStream>(std::move(file)));
+  };
+}
+
+// Removes unwanted characters for readability and to eliminate issues when
+// saving a file.
+std::string Sanitize(absl::string_view string) {
+  static const auto& kUnwantedChars = *new absl::flat_hash_set<char>{
+      '<', '>', ':', '\"', '/', '\\', '|', '?', '*', ' ', '(', ')'};
+
+  std::string sanitized;
+  sanitized.reserve(string.size());
+
+  bool skip = false;
+  for (const char& c : string) {
+    if (auto it = kUnwantedChars.find(c); it != kUnwantedChars.end()) {
+      skip = true;
+      continue;
+    }
+    if (skip) {
+      skip = false;
+      sanitized.push_back('_');
+    }
+    sanitized.push_back(c);
+  }
+
+  return sanitized;
+}
+
+// Pass instrumentation that dumps MLIR based on the criteria specified by
+// `ir_dump_*` debug options.
+//
+// While `mlir::PassManager::enableIRPrinting` provides a similar functionality,
+// it is cumbersome to manually copy printed IRs and run them with `litert-opt`.
+// Also, long MLIR dumps are often truncated during printing. Instead, this
+// instrumentation dumps MLIR to external directories for convenience.
+class DumpInstrumentation : public mlir::PassInstrumentation {
+ public:
+  explicit DumpInstrumentation(absl::string_view dump_dir,
+                               absl::string_view dump_pass_regex,
+                               absl::string_view dump_func_regex)
+      : dump_dir_(dump_dir),
+        dump_pass_re_(std::make_unique<RE2>(dump_pass_regex)),
+        dump_func_re_(std::make_unique<RE2>(dump_func_regex)) {}
+
+  DumpInstrumentation(const DumpInstrumentation& other) = delete;
+  DumpInstrumentation& operator=(const DumpInstrumentation& other) = delete;
+
+  void runBeforePass(mlir::Pass* pass, mlir::Operation* op) override {
+    // Always print before the first pass.
+    if (!printed_) {
+      Dump("before_all", op);
+      printed_ = true;
+    }
+
+    if (RE2::FullMatch(pass->getName(), *dump_pass_re_)) {
+      Dump(absl::StrCat(absl::string_view(pass->getName()), "_before"), op,
+           absl::StrCat(absl::Hex(pass_counter_, absl::kZeroPad8)));
+    }
+  }
+
+  void runAfterPass(mlir::Pass* pass, mlir::Operation* op) override {
+    if (RE2::FullMatch(pass->getName(), *dump_pass_re_)) {
+      Dump(absl::StrCat(absl::string_view(pass->getName()), "_after"), op,
+           absl::StrCat(absl::Hex(pass_counter_++, absl::kZeroPad8)));
+    }
+  }
+
+ private:
+  // Dumps the given op. `name` is used as part of the filename to help
+  // distinguish dumps at different passes.
+  void Dump(absl::string_view name, mlir::Operation* op,
+            std::string prefix = "") {
+    static constexpr char kFiletypeSuffix[] = "mlir";
+    // Find names of all func ops with public visibility and check whether at
+    // least one of them matches `dump_func_re_`.
+    llvm::SmallVector<absl::string_view> func_names;
+    bool match = false;
+    op->walk([&](mlir::func::FuncOp func) {
+      if (func.isPublic()) {
+        const absl::string_view name = func.getSymName();
+        if (name.empty()) {
+          return;
+        }
+
+        func_names.push_back(name);
+        if (RE2::FullMatch(name, *dump_func_re_)) {
+          match = true;
+        }
+      }
+    });
+    if (!func_names.empty() && !match) {
+      return;
+    }
+
+    // Sort the function names for determinism.
+    llvm::sort(func_names);
+
+    std::string joined_func_names = Sanitize(absl::StrJoin(func_names, "-"));
+    std::string sanitized_name = Sanitize(name);
+
+    std::vector<absl::string_view> name_parts;
+    if (!prefix.empty()) {
+      name_parts.emplace_back(prefix);
+    }
+    if (!joined_func_names.empty()) {
+      name_parts.emplace_back(joined_func_names);
+    }
+    name_parts.emplace_back(sanitized_name);
+    name_parts.emplace_back(kFiletypeSuffix);
+
+    // Build a filename such that it contains function names and pass names for
+    // easy disambiguation.
+    const std::string filename = tsl::io::JoinPath(
+        dump_dir_, absl::StrJoin(name_parts.begin(), name_parts.end(), "."));
+
+    // Open the file for dumping. Failures are logged instead of being
+    // propagated to the client because they are non-fatal.
+
+    std::unique_ptr<tsl::WritableFile> file;
+    if (auto status = tsl::Env::Default()->NewWritableFile(filename, &file);
+        !status.ok()) {
+      LOG(ERROR) << "Unable to open '" << filename
+                 << "' for dumping TFLite MLIR output: " << status;
+      return;
+    }
+    file = std::make_unique<tsl::BufferedWritableFile>(std::move(file));
+
+    WritableFileRawStream os(std::move(file));
+    op->print(os);
+  }
+
+  const std::string dump_dir_;
+  const std::unique_ptr<RE2> dump_pass_re_;
+  const std::unique_ptr<RE2> dump_func_re_;
+
+  // Counter used for pass name prefix to signify sequence
+  int pass_counter_ = 0;
+
+  bool printed_ = false;
+};
+
+std::function<bool(mlir::Pass*, mlir::Operation*)> CreatePrintIRFun(
+    const std::string& pass_regex) {
+  std::function<bool(mlir::Pass*, mlir::Operation*)> fun;
+  if (pass_regex.empty()) {
+    return fun;
+  }
+  return [pr = pass_regex](mlir::Pass* p, mlir::Operation*) {
+    static const RE2* const re = new RE2(pr);
+    if (RE2::FullMatch(p->getName(), *re)) {
+      return true;
+    }
+    return false;
+  };
+}
+
+}  // namespace
+
+void InitPassManager(mlir::PassManager& pm,
+                     const converter::DebugOptions& options,
+                     llvm::raw_ostream& out) {
+  std::string dump_dir = options.ir_dump_dir();
+
+  bool dump_to_dir = !dump_dir.empty();
+  bool print_to_stdout =
+      !options.print_ir_before().empty() || !options.print_ir_after().empty();
+
+  if (dump_to_dir || print_to_stdout) {
+    // Necessary for maintaining sequence of passes when dumping MLIR to files
+    // or stdout.
+    pm.getContext()->disableMultithreading();
+  }
+
+  if (dump_to_dir) {
+    dump_dir = tsl::io::JoinPath(
+        dump_dir, absl::FormatTime("%E4Y%m%d_%H%M%E6S", absl::Now(),
+                                   absl::LocalTimeZone()));
+
+    // Get a valid file path to dump with.
+    tsl::Env* env = tsl::Env::Default();
+    if (auto status = env->RecursivelyCreateDir(dump_dir); !status.ok()) {
+      LOG(WARNING) << "Failed to create '" << dump_dir
+                   << "' directory for dumping: " << status;
+      return;
+    }
+
+    // Set a default crash reproducer for easier debugging.
+    if (auto reproducer_stream_factory = GetReproducerStreamFactory(dump_dir)) {
+      pm.enableCrashReproducerGeneration(std::move(reproducer_stream_factory));
+    }
+
+    pm.addInstrumentation(std::make_unique<DumpInstrumentation>(
+        dump_dir, options.ir_dump_pass_regex(), options.ir_dump_func_regex()));
+  }
+
+  if (print_to_stdout) {
+    std::function<bool(mlir::Pass*, mlir::Operation*)>
+        should_print_ir_before_pass(
+            CreatePrintIRFun(options.print_ir_before()));
+    std::function<bool(mlir::Pass*, mlir::Operation*)>
+        should_print_ir_after_pass(CreatePrintIRFun(options.print_ir_after()));
+
+    mlir::OpPrintingFlags opPrintingFlags = mlir::OpPrintingFlags();
+
+    if (options.has_elide_elementsattrs_if_larger()) {
+      opPrintingFlags.elideLargeElementsAttrs(
+          options.elide_elementsattrs_if_larger());
+    }
+
+    pm.enableIRPrinting(should_print_ir_before_pass, should_print_ir_after_pass,
+                        options.print_ir_module_scope(),
+                        /*printAfterOnlyOnChange=*/true,
+                        /*printAfterOnlyOnFailure=*/false, out,
+                        opPrintingFlags);
+  }
+
+  // Enable pass timing. Note: MLIR expects `mlir::PassManager::enableTiming` to
+  // be called after all instrumentations are added.
+  if (options.enable_timing()) {
+    pm.enableTiming();
+  }
+
+  pm.addInstrumentation(
+      std::make_unique<mlir::TFL::ErrorCollectorInstrumentation>(
+          pm.getContext()));
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/lite/debug/debug.h

@@ -0,0 +1,34 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_DEBUG_DEBUG_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_DEBUG_DEBUG_H_
+
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"
+
+namespace tensorflow {
+
+// Initializes the pass manager with default options that make debugging easier.
+// The `out` method parameter is exposed for testing purposes and not intended
+// to be specified by client code.
+void InitPassManager(mlir::PassManager& pm,
+                     const converter::DebugOptions& options,
+                     llvm::raw_ostream& out = llvm::outs());
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_DEBUG_DEBUG_H_

tensorflow/compiler/mlir/lite/debug/debug_options.proto

@@ -0,0 +1,61 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+syntax = "proto2";
+
+package tensorflow.converter;
+
+// Additional parameters that control the debug behavior of the Converter.
+//
+// Next ID: 9
+message DebugOptions {
+  // If not empty, dumps MLIR to the specified directory. The initial state of
+  // the MLIR after import will be dumped at the beginning of each pass manager
+  // run. Additionally, MLIR will be dumped before and after each pass
+  // depending on pass names and functions matched using the
+  // `ir_dump_pass_regex` and `ir_dump_func_regex` values.
+  optional string ir_dump_dir = 1 [default = ""];
+
+  // Regular expression that matches the names of passes in pascal case (e.g.,
+  // FooPass) before/after which MLIR will be dumped. Effective only if
+  // ir_dump_dir is not empty.
+  optional string ir_dump_pass_regex = 2 [default = ".*"];
+
+  // Regular expression that matches the names of functions to be dumped. MLIR
+  // modules are dumped only if there's at least one public function in the
+  // module whose name matches the pattern. Effective only if ir_dump_dir is
+  // not empty.
+  optional string ir_dump_func_regex = 3 [default = ".*"];
+
+  // If true, report the execution time of each MLIR pass.
+  optional bool enable_timing = 4 [default = false];
+
+  // Prints MLIR before specified passes. Supports regular expressions for
+  // matching against the names of the desired passes.
+  optional string print_ir_before = 5 [default = ""];
+
+  // Prints MLIR after specified passes. Supports regular expressions for
+  // matching against the names of the desired passes. Currently only prints
+  // after a pass if the MLIR is mutated.
+  optional string print_ir_after = 6 [default = ""];
+
+  // If true, always print the top-level operation when printing IR for
+  // print_ir_[before|after].
+  optional bool print_ir_module_scope = 7 [default = true];
+
+  // Elide ElementsAttrs with \"...\" that have more elements than the given
+  // upper limit.
+  optional int64 elide_elementsattrs_if_larger = 8;
+}

tensorflow/compiler/mlir/lite/debug/debug_test.cc

@@ -0,0 +1,288 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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 "tensorflow/compiler/mlir/lite/debug/debug.h"
+
+#include <stdint.h>
+
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinDialect.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/env.h"
+#include "tensorflow/core/platform/types.h"
+#include "tsl/platform/path.h"
+
+namespace tensorflow {
+namespace debug_test {
+
+class NopPass : public mlir::PassWrapper<NopPass, mlir::OperationPass<>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(NopPass)
+
+  void runOnOperation() override {}
+};
+
+class MutatePass : public mlir::PassWrapper<MutatePass, mlir::OperationPass<>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MutatePass)
+
+  void runOnOperation() override {
+    mlir::OpBuilder builder(&getContext());
+    getOperation()->setAttr("tfl.random_attr", builder.getUnitAttr());
+  }
+};
+
+class AlwaysFailPass
+    : public mlir::PassWrapper<AlwaysFailPass, mlir::OperationPass<>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AlwaysFailPass)
+
+  void runOnOperation() override { signalPassFailure(); }
+};
+
+}  // namespace debug_test
+
+namespace {
+
+using ::testing::HasSubstr;
+using ::testing::IsEmpty;
+using ::testing::Not;
+using namespace tensorflow::debug_test;
+
+class InitPassManagerTest : public testing::Test {
+ protected:
+  InitPassManagerTest()
+      : path_(GetOutputPath()), context_([]() {
+          mlir::registerPassManagerCLOptions();
+          mlir::DialectRegistry registry;
+          registry.insert<mlir::BuiltinDialect>();
+          registry.insert<mlir::arith::ArithDialect>();
+          registry.insert<mlir::func::FuncDialect>();
+          registry.insert<mlir::TFL::TensorFlowLiteDialect>();
+          return registry;
+        }()) {
+    context_.loadAllAvailableDialects();
+
+    mlir::OpBuilder builder(&context_);
+    module_ = mlir::ModuleOp::create(builder, builder.getUnknownLoc());
+
+    builder.setInsertionPointToStart(module_->getBody());
+    auto func = mlir::func::FuncOp::create(builder,  //
+                                           builder.getUnknownLoc(), "main",
+                                           builder.getFunctionType({}, {}));
+    func->setAttr("tfl.func", builder.getUnitAttr());
+    builder.setInsertionPointToStart(func.addEntryBlock());
+    llvm::SmallVector<int> shape{1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
+    mlir::arith::ConstantOp::create(
+        builder, builder.getUnknownLoc(),
+        mlir::DenseIntElementsAttr::get(
+            mlir::RankedTensorType::get(shape.size(), builder.getI32Type()),
+            shape));
+    mlir::func::ReturnOp::create(builder, builder.getUnknownLoc());
+  }
+
+  absl::Status GetDumpDir(std::string* dump_dir) {
+    std::vector<std::string> files;
+    if (auto status = tsl::Env::Default()->GetChildren(path_, &files);
+        !status.ok()) {
+      return status;
+    }
+    if (files.size() != 1) {
+      return absl::FailedPreconditionError(
+          "Expecting directory to have one child.");
+    }
+    *dump_dir = tsl::io::JoinPath(path_, files[0]);
+    return absl::OkStatus();
+  }
+
+  std::string path_;
+  mlir::MLIRContext context_;
+  mlir::OwningOpRef<mlir::ModuleOp> module_;
+
+ private:
+  std::string GetOutputPath() {
+    const auto* const test_info =
+        testing::UnitTest::GetInstance()->current_test_info();
+    return tsl::io::JoinPath(
+        getenv("TEST_UNDECLARED_OUTPUTS_DIR"),
+        absl::StrCat(test_info->test_suite_name(), ".", test_info->name()));
+  }
+};
+
+TEST_F(InitPassManagerTest, CrashReproducer) {
+  converter::DebugOptions debug_options;
+  *debug_options.mutable_ir_dump_dir() = path_;
+
+  mlir::PassManager pm(&context_);
+  InitPassManager(pm, debug_options);
+  pm.addPass(std::make_unique<AlwaysFailPass>());
+  ASSERT_TRUE(mlir::failed(pm.run(*module_)));
+
+  std::string dump_dir;
+  TF_ASSERT_OK(GetDumpDir(&dump_dir));
+
+  std::string mlir_dump;
+  TF_ASSERT_OK(tsl::ReadFileToString(
+      tsl::Env::Default(),
+      tsl::io::JoinPath(dump_dir, "tfl_mlir_crash_repro.mlir"), &mlir_dump));
+  EXPECT_THAT(mlir_dump, Not(IsEmpty()));
+}
+
+TEST_F(InitPassManagerTest, DumpToDir) {
+  converter::DebugOptions debug_options;
+  *debug_options.mutable_ir_dump_dir() = path_;
+  *debug_options.mutable_ir_dump_pass_regex() = R"(.*NopPass)";
+
+  mlir::PassManager pm(&context_);
+  InitPassManager(pm, debug_options);
+  pm.addPass(std::make_unique<NopPass>());
+  ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
+  std::string dump_dir;
+  TF_ASSERT_OK(GetDumpDir(&dump_dir));
+
+  {
+    std::string mlir_dump;
+    TF_ASSERT_OK(tsl::ReadFileToString(
+        tsl::Env::Default(),
+        tsl::io::JoinPath(
+            dump_dir, "00000000.main.tensorflow_debug_test_NopPass_after.mlir"),
+        &mlir_dump));
+    EXPECT_THAT(mlir_dump, Not(IsEmpty()));
+  }
+  {
+    std::string mlir_dump;
+    TF_ASSERT_OK(tsl::ReadFileToString(
+        tsl::Env::Default(),
+        tsl::io::JoinPath(
+            dump_dir,
+            "00000000.main.tensorflow_debug_test_NopPass_before.mlir"),
+        &mlir_dump));
+    EXPECT_THAT(mlir_dump, Not(IsEmpty()));
+  }
+}
+
+TEST_F(InitPassManagerTest, PrintIRBeforeEverything) {
+  converter::DebugOptions debug_options;
+  *debug_options.mutable_print_ir_before() = R"(.*)";
+  std::string captured_out;
+  llvm::raw_string_ostream out(captured_out);
+
+  mlir::PassManager pm(&context_);
+  InitPassManager(pm, debug_options, out);
+  pm.addPass(std::make_unique<NopPass>());
+  ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
+
+  EXPECT_THAT(captured_out,
+              HasSubstr("IR Dump Before tensorflow::debug_test::NopPass"));
+  EXPECT_THAT(captured_out,
+              Not(HasSubstr("IR Dump After tensorflow::debug_test::NopPass")));
+}
+
+TEST_F(InitPassManagerTest, PrintIRAfterEverything) {
+  converter::DebugOptions debug_options;
+  *debug_options.mutable_print_ir_after() = R"(.*)";
+  std::string captured_out;
+  llvm::raw_string_ostream out(captured_out);
+
+  mlir::PassManager pm(&context_);
+  InitPassManager(pm, debug_options, out);
+  pm.addPass(std::make_unique<MutatePass>());
+  ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
+
+  EXPECT_THAT(captured_out,
+              HasSubstr("IR Dump After tensorflow::debug_test::MutatePass"));
+  EXPECT_THAT(
+      captured_out,
+      Not(HasSubstr("IR Dump Before tensorflow::debug_test::MutatePass")));
+}
+
+TEST_F(InitPassManagerTest, PrintIRBeforeAndAfterEverything) {
+  converter::DebugOptions debug_options;
+  *debug_options.mutable_print_ir_before() = R"(.*)";
+  *debug_options.mutable_print_ir_after() = R"(.*)";
+  std::string captured_out;
+  llvm::raw_string_ostream out(captured_out);
+
+  mlir::PassManager pm(&context_);
+  InitPassManager(pm, debug_options, out);
+  pm.addPass(std::make_unique<MutatePass>());
+  ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
+
+  EXPECT_THAT(captured_out,
+              HasSubstr("IR Dump After tensorflow::debug_test::MutatePass"));
+  EXPECT_THAT(captured_out,
+              HasSubstr("IR Dump Before tensorflow::debug_test::MutatePass"));
+}
+
+TEST_F(InitPassManagerTest, ElideLargeElementAttrs) {
+  converter::DebugOptions debug_options;
+  *debug_options.mutable_print_ir_before() = R"(.*)";
+  debug_options.set_elide_elementsattrs_if_larger(5);
+  std::string captured_out;
+  llvm::raw_string_ostream out(captured_out);
+
+  mlir::PassManager pm(&context_);
+  InitPassManager(pm, debug_options, out);
+  pm.addPass(std::make_unique<MutatePass>());
+  ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
+
+  EXPECT_THAT(captured_out, HasSubstr("dense_resource<__elided__>"));
+}
+
+TEST_F(InitPassManagerTest, DontElideSmallerElementAttrs) {
+  converter::DebugOptions debug_options;
+  *debug_options.mutable_print_ir_before() = R"(.*)";
+  debug_options.set_elide_elementsattrs_if_larger(11);
+  std::string captured_out;
+  llvm::raw_string_ostream out(captured_out);
+
+  mlir::PassManager pm(&context_);
+  InitPassManager(pm, debug_options, out);
+  pm.addPass(std::make_unique<MutatePass>());
+  ASSERT_TRUE(mlir::succeeded(pm.run(*module_)));
+
+  EXPECT_THAT(captured_out,
+              HasSubstr("dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]>"));
+}
+
+}  // namespace
+}  // namespace tensorflow