namespace {

// Retrieve the MLIR dump directory. The directory is read from the environment
// variable `TF_QUANT_MLIR_DUMP_PREFIX`. However, if a special value "sponge" is
// set to `TF_QUANT_MLIR_DUMP_PREFIX`, it uses the directory set in
// `TEST_UNDECLARED_OUTPUT_DIRS`. Returns `absl::FailedPreconditionError` if
// either:
//   1. `TF_QUANT_MLIR_DUMP_PREFIX` is not set (empty), or
//   2. `TEST_UNDECLARED_OUTPUT_DIRS` is not set (empty) when

                      absl::string_view file_name_prefix);

// If verbosity level >= 1, this will dump intermediate IRs of passes to a file.
// The dumped mlir files with be under a directory determined by
// the TF_QUANT_MLIR_DUMP_PREFIX env variable. The PassManager will dump to a
// new file for each pass. The file name will have the format
// {file_name_prefix}_{pass_number}_{pass_name}_{before|after}.mlir.
// * `file_name_prefix` is from input.

    if (!tsl::io::GetTestUndeclaredOutputsDir(&test_dir_)) {
      test_dir_ = tsl::testing::TmpDir();
    }
  }

  void SetUp() override {
    tsl::setenv("TF_QUANT_MLIR_DUMP_PREFIX", test_dir_.c_str(), 1);

    mlir::DialectRegistry dialects;
    dialects.insert<mlir::BuiltinDialect, mlir::func::FuncDialect,
                    mlir::stablehlo::StablehloDialect>();