// StableHLO -> MHLO legalization for MHLO optimization.
  pm.addPass(mlir::mhlo::createStablehloLegalizeToHloPass());
  // Rewrites legacy StableHLO ops.
  AddUnfuseMhloOpsPasses(pm);
  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
  // MHLO -> StableHLO legalization.
  pm.addPass(mlir::mhlo::createHloLegalizeToStablehloPass());
}

}

def LegalizeStablehloToVhloPass : Pass<"stablehlo-legalize-vhlo", "ModuleOp"> {
  let summary = "Legalize StableHLO to VHLO for ODML.";
  let dependentDialects = ["mlir::vhlo::VhloDialect"];
}

def LegalizeVhloToStablehloPass : Pass<"vhlo-legalize-stablehlo", "ModuleOp"> {
  let summary = "Legalize VHLO to StableHLO for ODML.";
  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
}

// The TF dialect uses some TF types that are illegal in the MHLO dialect and
// some generic types that are legal in MHLO. This pass legalizes TF types into
// types that are legal in MHLO. For example, TF::Qint8Type is converted to i8.
// Rewrites here should run before TF to MHLO op legalizations are run.

#include <memory>

#include "absl/log/log.h"
#include "llvm/ADT/STLExtras.h"

#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/client/xla_builder.h"
#include "xla/client/xla_computation.h"
#include "xla/hlo/ir/hlo_instruction.h"
#include "xla/hlo/ir/hlo_opcode.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/service/hlo.pb.h"
#include "xla/translate/hlo_to_mhlo/hlo_function_importer.h"
#include "xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h"

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace mhlo {

ConversionTarget GetDefaultLegalConversionTargets(MLIRContext& mlir_context,
                                                  bool legalize_chlo) {
  ConversionTarget target(mlir_context);

  if (legalize_chlo) {
    target.addIllegalDialect<chlo::ChloDialect>();

    }
    Value new_const_op =
        rewriter.create<mhlo::ConstantOp>(bcast_op.getLoc(), result);
    rewriter.replaceOp(bcast_op, {new_const_op});
    return success();
  }
};

using FoldBroadcastInDimBeforeMulOp =
    FoldBroadcastInDimBeforeBinaryElementwiseOp<mhlo::MulOp>;

// Constant folds mhlo.mul, this folder doesn't have an upper limit on how many
// elements can be folded.

#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"
#include "xla/mlir_hlo/mhlo/transforms/rewriters.h"

//===----------------------------------------------------------------------===//
// The unfuse-mhlo-batch-norm Pass.
//===----------------------------------------------------------------------===//

namespace mlir::quant::stablehlo {

  ConversionTarget target(context);
  target.addLegalDialect<TFL::TensorFlowLiteDialect, mhlo::MhloDialect>();
  target.addLegalOp<func::CallOp, func::ConstantOp, arith::ConstantOp>();
  target.addDynamicallyLegalOp<mhlo::CustomCallOp>(IsCustomCallLegal);
  target.addDynamicallyLegalOp<mhlo::ReduceOp>(IsReduceOpLegal);
  // Converted MHLO ops should be marked illegal here.
  // TODO: b/304003568 - Add TF_TransposeOp folding logic to tflite.

  if (lower_to_xla_hlo) {
    // This pass operates on MHLO control flow ops so it should be legalized
    // after the control flow ops are legalized.
    pm.addPass(mlir::mhlo::CreateLegalizeTFCommunicationPass());

    // Everything should be MHLO after this.
    if (!allow_partial_conversion) {
      pm.addNestedPass<mlir::func::FuncOp>(
          mlir::mhlo::CreateVerifyTFXLALegalizationPass(legalize_chlo));
    }
  }

  ];
  let dependentDialects = [
    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",
    "mlir::func::FuncDialect", "mlir::mhlo::MhloDialect",
    "mlir::quant::QuantizationDialect", "mlir::chlo::ChloDialect",
    "mlir::vhlo::VhloDialect", "mlir::shape::ShapeDialect",
    "mlir::quantfork::QuantizationForkDialect",
  ];
}