// This is the legalization pattern definition file for HLO to TF.

include "mlir/IR/OpBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "stablehlo/dialect/ChloOps.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mhlo/IR/hlo_ops.td"

// Check if broadcasting is compatible with TF ops.
def IsLegalNumpyRankedBroadcast :

#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/local_client.h"
#include "xla/executable_run_options.h"
#include "xla/hlo/ir/hlo_input_output_alias_config.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/tf_pjrt_client.h"
#include "xla/service/executable.h"
#include "xla/service/gpu/gpu_executable_run_options.h"

  // have `xla_device_metadata`. This is used for TFRT-TPU when
  // BuildXlaDeviceCompiler() is called in GetCompilerIr(). Currently only
  // lowering to HLO is needed there and xla::LocalClient doesn't support
  // building the executable for TFRT-TPU and hence, is set to nullptr here.
  if (platform_info.device_type() == DEVICE_TPU) {
    *xla_device_compiler = CreateXlaDeviceCompiler(

  // Convert model.
  if (enable_mlir_converter) {
    if (model_flags.use_hlo_import() && model_flags.has_saved_model_dir()) {
      PyErr_SetString(PyExc_ValueError,
                      "Cannot specify both saved_model and hlo import.");
      return nullptr;
    }

    if (model_flags.use_hlo_import()) {
      status = tensorflow::ConvertJaxToTFLiteFlatBuffer(
          input_contents_txt, model_flags, toco_flags,

        "@llvm-project//mlir:Support",
        "@local_xla//xla:status_macros",
        "@local_xla//xla:xla_data_proto_cc",
        "@local_xla//xla/client:sharding_builder",
        "@local_xla//xla/hlo/ir:hlo",
        "@local_xla//xla/service:hlo_parser",
    ],
)

cc_library(
    name = "location_utils",
    srcs = ["utils/location_utils.cc"],
    hdrs = ["utils/location_utils.h"],
    deps = [

limitations under the License.
==============================================================================*/

// This file implements logic for lowering TensorFlow dialect's communication
// ops (TF/XLA) to the HLO dialect.

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"

      Env::Default(), GetFilePath(key2, persistor.persistent_cache_directory()),
      entry));

  MockXlaCompilerClient mock_client;
  // Try to load executable corresponding to key2 (whose file contains HLO
  // corresponding to key1).
  auto loaded_executable = persistor.TryToLoadExecutable(
      /*signature_hash=*/123, "signature", DefaultXlaOptions(),
      compilation_result_mul, &mock_client);

      TypeID::get<TF::ConstOp>(),

      // AssertOp with string types are not supported by the fallback.
      TypeID::get<TF::AssertOp>(),

      // TF2XLA fallback pattern doesn't support these op as MLIR hlo builder
      // doesn't override the necessary builder methods. These ops have simple
      // lowering pattern so this should be safe.
      TypeID::get<TF::CrossReplicaSumOp>(),
      TypeID::get<TF::InfeedDequeueTupleOp>(),

  AddCanonicalizationPatterns(module.getContext(), &patterns);

  // `supported_ops` contains the name of all of the ops that can potentially be
  // lowered into HLO on the device. This doesn't always mean that the op can
  // be lowered in the future passes but if the op is not in this set, it can't
  // be lowered in a subsequent pass.
  llvm::DenseSet<OperationName> supported_ops;

// COM: This file is there to check that the `tfl-legalize-hlo` pass exists in `odml-to-stablehlo-opt`.

// RUN: odml-to-stablehlo-opt %s -tfl-legalize-hlo -split-input-file | FileCheck %s --dump-input=fail

func.func @main(%arg0: tensor<5x7xf32>) -> tensor<5x7xf32> {
  func.return %arg0: tensor<5x7xf32>
// CHECK-LABEL: main
// CHECK: return %arg0 : tensor<5x7xf32>
}

// - transpose
//