}
  })";

TEST(LegalizeMlirTest, LegalizesModule) {
  mlir::DialectRegistry mlir_registry;
  RegisterAllTensorFlowDialects(mlir_registry);

  std::vector<tensorflow::TensorShape> arg_shapes;
  XlaCompilationResult compilation_result;
  auto status = CompileSerializedMlirToXlaHlo(
      kMlirModuleStr, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
      /*use_tuple_args=*/true, /*enable_op_fallback=*/false,

    } break;
    default:
      VLOG(2) << "LegalizeMlirToHlo computation input: unknown";
      break;
  }
}

Status DumpHloCompilationResult(std::string_view name,
                                XlaCompilationResult* compilation_result) {
  if (!VLOG_IS_ON(2) &&
      !DEBUG_DATA_DUMPER()->ShouldDump(std::string(name), kDebugGroupMain)) {
    return absl::OkStatus();
  }

  TF_ASSIGN_OR_RETURN(

                      llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
                          custom_legalization_passes);

// Apply shape, description, and resource information to inputs and outputs
// in the XlaCompilationResult. This should be called after
// compilation_result->computation was set.
ABSL_DEPRECATED("Not meant to be used directly and should be a util.")
Status PopulateResultIOInfo(

namespace tensorflow {
namespace tf2xla {
namespace internal {

// Legalize the given MLIR module to XLA HLO using a combination of the MLIR
// Bridge and XlaBuilder
absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,

    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes);

// Compiles a serialized MLIR module into XLA HLO, generates all accompanying
// metadata and stores them in CompilationResult.
absl::StatusOr<XlaCompilationResult> LegalizeWithMlirBridge(
    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,

namespace tensorflow {
namespace tf2xla {
namespace internal {

using metrics::IncrementTfMlirBridgeSecondPhaseCounter;
using metrics::MlirBridgeSecondPhaseMetric;
using tpu::MlirToHloArgs;

absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,

  std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
  std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;

  auto compilation_result = std::make_unique<XlaCompilationResult>();

  return CompileFromMlirToXlaHlo(
      compile_to_xla_hlo, mlir_to_hlo_args, metadata_proto,
      /*device_type=*/"XLA_TPU_JIT",
      /*shape_determination_fns=*/{}, use_tuple_args, compilation_result.get(),

  CellReader<int64_t> reader(kMetric);

  ASSERT_THAT(failed(), Not(IncrementedOrFiltered(reader.Delta(kOkStatus), 1)));
}

tensorflow::XlaCompilationResult CreateXlaComputationResult(
    const char* hlo_name) {
  auto result = tensorflow::XlaCompilationResult();
  xla::HloModuleProto hlo;
  hlo.set_name(hlo_name);
  result.computation = std::make_shared<xla::XlaComputation>(hlo);
  return result;
}

                           arg_core_mapping, per_core_arg_shapes);
  if (!sharding_result.ok()) {
    return sharding_result;
  }
  return compiled_mlir;
}

absl::StatusOr<XlaCompilationResult> LegalizeWithMlirBridge(
    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,

//  arg_core_mapping - Which args go on which cores.
//  per_core_arg_shapes - For each core, the shapes for each argument.
//  client - The Xla Compilation client.
absl::StatusOr<tensorflow::XlaCompilationResult> LegalizeMlirToHlo(
    const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,