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

// 2) Use fallback with a device that supports all ops (ops can be legalized).
// 3) Use fallback with unspecified device (ops cannot be legalized).
// 4) Use fallback with specified but unsupported device (ops cannot be legalized).
//
// Note: For 3) and 4) we do not use `-verify-diagnostics` because these cases
// produce remarks that don't occur for 1) and 2) and there is no way to check

#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,

def LegalizeTFCommunicationPass : Pass<"xla-legalize-tf-communication", "ModuleOp"> {
  let summary = "Legalize TF/XLA communication ops (TensorFlow dialect) to the HLO "
           "dialect";
  let dependentDialects = ["mhlo::MhloDialect", "sparse_tensor::SparseTensorDialect"];
  let constructor = "::mlir::mhlo::CreateLegalizeTFCommunicationPass()";
  let description = [{
    A pass that legalizes TF/XLA communication ops, propagates their respective

    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();
    }

#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project

namespace mlir {
namespace mhlo {

// Returns a ConversionTarget that includes default legalized MLIR dialects
// for conversion to XLA.
// If legalize_chlo is true, the resulting conversion target cannot have CHLO.
mlir::ConversionTarget GetDefaultLegalConversionTargets(
    MLIRContext& mlir_context, bool legalize_chlo);

// 1) Reshape/transpose are inserted because tfl.BatchMatMul requires
// size(contracting_dimensions) = 1 and size(output_dim) = 1, whereas
// mhlo.dot_general has no such restriction.
// 2) Inserted mhlo.reshape/transpose will be legalized to tf.reshape/transpose
// in LegalizeHloToTf (then from tf to tfl later).
// 3) If the operands are dynamic shaped tensors, mhlo.DynamicReshapeOp is
// inserted instead of the regular reshape, and additional ops (e.g. Gather,

      ConversionPatternRewriter& rewriter) const final;
};

// Ops that have a call_target_name starting with the prefix "custom_call." and
// backend_config of type StringAttr (if specified) should be legalized (i.e.
// considered to be illegal if still present after this pass is complete).
std::optional<bool> IsCustomCallLegal(mhlo::CustomCallOp op);

}  // namespace odml
}  // namespace mlir