module_name.str(), dump_group_name, pass_tag_name);
      },
      /*pass_prefix=*/"",
      /*print_module_scope=*/true));
  pm.enableTiming();
}

}  // namespace

void AddTPULowerClusterToRuntimeOpsPassPipeline(OpPassManager& pm,
                                                llvm::StringRef module_name) {
  pm.addPass(mlir::TFTPU::CreateTPURewritePass(module_name));

      /*filter_stack=*/!VLOG_IS_ON(1));

  if (VLOG_IS_ON(1) ||
      DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
    ::tensorflow::DumpMlirOpToFile(DEBUG_DATA_DUMPER()->GetDumpFilename(
                                       module_name.str(), kDebugGroupMain,
                                       "_" + dump_prefix.str() + "_before"),

      DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
    ::tensorflow::DumpMlirOpToFile(
        DEBUG_DATA_DUMPER()->GetDumpFilename(module_name.str(), kDebugGroupMain,
                                             dump_prefix.str() + "_before"),
        module, llvm::StringRef(), &bridge);
  }

  if (VLOG_IS_ON(2) ||
      DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(),

            module_name.str(), kDebugGroupMain,
            "tfxla_bridge_v2_tfdialect_to_executor_before"),
        module, llvm::StringRef(), &tf_to_executor);

    if (VLOG_IS_ON(2) ||
        DEBUG_DATA_DUMPER()->ShouldDump(
            module_name.str(), kDebugGroupBridgePhase1ExecutorExport)) {
      internal::EnablePassIRPrinting(
          tf_to_executor, kDebugGroupBridgePhase1ExecutorExport, module_name);
    }
  }

      module.getContext(), /*propagate=*/false,
      /*filter_stack=*/!VLOG_IS_ON(1));

  if (VLOG_IS_ON(1) ||
      DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) {
    ::tensorflow::DumpMlirOpToFile(
        DEBUG_DATA_DUMPER()->GetDumpFilename(
            module_name.str(), kDebugGroupMain,
            "tfxla_bridge_v1_tfdialect_to_executor_before"),
        module, llvm::StringRef(), &tf_to_executor);

struct TPURewritePass : public impl::TPURewritePassBase<TPURewritePass> {
  explicit TPURewritePass(llvm::StringRef _module_name)
      : module_name(_module_name) {}

  void runOnOperation() override;

  llvm::StringRef module_name;
};

// Creates a missing attribute error message.
std::string CreateMissingAttributeMsg(llvm::StringRef attribute) {

  for (const auto& creator : pass_instrumentors) {
    tf2xla.addInstrumentation(creator());
  }
  if (DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain) ||
      VLOG_IS_ON(1)) {
    tensorflow::DumpMlirOpToFile(
        DEBUG_DATA_DUMPER()->GetDumpFilename(module_name.str(), kDebugGroupMain,
                                             "legalize_hlo_before"),
        module_op, "", &tf2xla);
  }

namespace tf2xla {
namespace internal {

// Given the pass manager, add Bridge passes to cluster the replicated input
// graphs.
void AddReplicatedBridgeClusteringPipelinePasses(
    mlir::OpPassManager& pm, llvm::StringRef module_name = llvm::StringRef());

// Same as above but for non replicated graphs.
void AddNonReplicatedBridgeClusteringPipelinePasses(mlir::OpPassManager& pm);

};  // namespace internal
};  // namespace tf2xla

//   replicated, e.g. CPU/GPU graphs. is_in_fallback_enabled_mode - Whether this
//   was called with fallback to the non-MLIR Bridge. This is just for logging
//   purposes and doesn't affect logic. module_name - What the input module name
//   is for debugging help.
//
// Output: Modifies the input module in place with clustered operations.
//   status - Whether the transformation to cluster the input MLIR module was

    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns =
        {},
    llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes = {},
    llvm::StringRef module_name = llvm::StringRef());

// Creates a MLIR pipeline that lowers MLIR module to MHLO dialect. The input
// module should only contain operations in tf dialect. For example, if the