include "mlir/Pass/PassBase.td"

def LegalizeTF : Pass<"xla-legalize-tf", "ModuleOp"> {
  let summary = "Legalize from TF dialect's or HLO dialect's control flow.";

  let description = [{
    Legalizes from TF dialect to HLO dialect. When allow_partial_conversion is
    false, emits an error if there is any operation that can't be legalized.
    When `tf2xla_fallback_device_type` is not `None`, also uses legalization

#include "mlir/Support/TypeID.h"  // from @llvm-project

namespace mlir {
namespace mhlo {

// Given the Operation, return True if the op is legalized with MLIR.
// Returns false if it is legalized with an old TF2XLA Kernel.
bool IsOpLegalizedWithMlir(Operation& op);

// Given the type ID, check if it's legalized with MLIR.
bool IsTypeLegalizedWithMlir(const TypeID& type_id);

// Returns true if the op is considered a dynamic padder op.

def VerifyQuantLegalization : Pass<"verify-quant-legalization", "mlir::func::FuncOp"> {
  let summary = "Verifies that all TF quant ops and types have been legalized.";

  let description = [{
    Ensures that all TF quant ops and types have been legalized to HLO
    and reports an error about which op failed to legalize. This pass
    does not transform any ops and is checking.}];

  let constructor = "mlir::quant::stablehlo::CreateVerifyQuantLegalizationPass()";

/// allow_partial_conversion is false, emits an error if there is any operation
/// that can't be legalized.
/// When `tf2xla_fallback_device_type` is not `None`, also uses legalization
/// patterns from TF2XLA fallback for provided device type (see
/// legalize_tf_with_tf2xla.cc for details). By default, TF2XLA fallback is not
/// used.
LogicalResult legalizeTF(
    Operation* op, bool allow_partial_conversion = false,
    bool legalize_chlo = true,

    if (IsSupportedTFCustomOp(op)) return WalkResult::advance();

    // Check for ops that are legalized to TFLite.
    if (op->getDialect()->getNamespace() == "tfl") {
      return WalkResult::advance();
    }
    // Check for ops that are not legalized to TFLite.
    if (IsSupportedTFDataForwardingOp(op)) {
      return WalkResult::advance();
    }

//   . Legalizes control flow operations.
//   . Decomposes compound resource operations so that the only remaining
//     operations on resource variables are resource reads/writes..
//   . Replaces resource reads/writes with function inputs/outputs and
//     eliminates the use of resource variables.
//   . Legalizes the operations to XLA HLO operations.

namespace mlir {
namespace TFL {
namespace {
#define GEN_PASS_DEF_LEGALIZEVARIABLESPASS
#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

// Attribute name to identify whether variables should be legalized to TFLite or
// not.
const char kLegalizeTflVariables[] = "tfl._legalize_tfl_variables";

bool HasSupportedElementType(Operation* op) {
  return utils::IsSupportedVariableType(op);
}

  target.addLegalDialect<func::FuncDialect>();
  target.addLegalDialect<tensor::TensorDialect>();
  target.addLegalDialect<shape::ShapeDialect>();
  target.addLegalOp<func::CallOp>();

  // These ops are legalized in LegalizeTFCommunication after this and that pass
  // only operates on MHLO control flow ops.
  target.addLegalOp<TF::_XlaHostComputeMlirOp, TF::XlaSendToHostOp,
                    TF::XlaRecvFromHostOp>();

      index_lt.getLhs() != body.getArgument(1) ||
      index_lt.getRhs() != body.getArgument(3))
    return failure();

  return success();
}

// Returns true if the given reduce op can be legalized to ArgMax/ArgMin ops.
std::optional<bool> IsReduceOpLegal(mhlo::ReduceOp reduce_op) {
  if (succeeded(MatchReduceToArgMinMaxType1(reduce_op, true, true)) ||
      succeeded(MatchReduceToArgMinMaxType1(reduce_op, false, true)) ||

class OperationPass;
class Type;

namespace TFL {

// Creates an instance of the TensorFlow Lite dialect LegalizeTF pass.
// When the given run_tfl_runtime_verification value is true, it will check each
// TFL builtin op towards the TFL runtime capability and the incompatible TF ops
// will be left in the graph without getting legalized. If `preserve_assert_op`
// is true, the TF::AssertOp will not be removed.