#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"

namespace tensorflow {

// Counter is used as a prefix for filenames.
static std::atomic<int> log_counter(0);

BridgeLoggerConfig::BridgeLoggerConfig(bool print_module_scope,
                                       bool print_after_only_on_change,
                                       mlir::OpPrintingFlags op_printing_flags)

  // inliner_pass should not.
  setenv("MLIR_BRIDGE_LOG_PASS_FILTER",
         "TPUResourceReadsWritesPartitioningPass;TensorFlowShapeInferencePass",
         1);
  BridgeLoggerConfig logger_config;
  EXPECT_TRUE(logger_config.ShouldPrint(partitioning_pass.get(),
                                        mlir_module_with_add.get()));
  EXPECT_TRUE(logger_config.ShouldPrint(shape_inference_pass.get(),

#include <functional>
#include <string>

#include "tensorflow/compiler/mlir/tensorflow/utils/bridge_logger.h"

namespace tensorflow {

class DataDumperLoggerConfig : public ::tensorflow::BridgeLoggerConfig {
 public:
  explicit DataDumperLoggerConfig(
      std::function<std::string(const std::string &, mlir::Operation *op)>
          get_filename,

        get_filename,
    const std::string &pass_prefix, bool print_module_scope,
    bool print_after_only_on_change, mlir::OpPrintingFlags op_printing_flags)
    : ::tensorflow::BridgeLoggerConfig(
          print_module_scope, print_after_only_on_change, op_printing_flags),
      get_filename_(get_filename),
      pass_prefix_(pass_prefix) {}

void DataDumperLoggerConfig::printBeforeIfEnabled(

// level >= 1, `bridge_logger.cc` prints details about pass invocations for
// which the IR dumping was skipped because of a filter.
class BridgeLoggerConfig : public mlir::PassManager::IRPrinterConfig {
 public:
  explicit BridgeLoggerConfig(
      bool print_module_scope = false, bool print_after_only_on_change = true,
      mlir::OpPrintingFlags op_printing_flags = mlir::OpPrintingFlags());

    // Print the whole module after each pass which requires disabling
    // multi-threading as well.
    module_op.getContext()->disableMultithreading();
    pm.enableIRPrinting(std::make_unique<tensorflow::BridgeLoggerConfig>(
        /*print_module_scope=*/true));
  }

  mlir::StatusScopedDiagnosticHandler diag_handler(module_op.getContext());
  if (failed(pm.run(module_op))) return diag_handler.ConsumeStatus();