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

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

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 351 : i32}} {

  ];
  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",
  ];
}

#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "mlir/IR/Location.h"  // from @llvm-project
#include "mlir/IR/Types.h"  // from @llvm-project
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace mhlo {

// Builds body for reduce op by using the template binary op as the
// reducer op.
template <typename Op>
void BuildReduceBody(Type element_type, Region* body, OpBuilder* builder) {

// RUN: tf-mlir-translate -mlir-tf-to-hlo-text %s -tf-input-shapes=128,10:10,1024:128,1024 -tf-xla-emit-use-tuple-args -tf-xla-emit-return-tuple | FileCheck %s

module attributes {tf.versions = {producer = 179 : i32}} {
  func.func @main(%arg0: tensor<128x10xf32> {mhlo.sharding = "\08\03\1A\02\01\02\22\02\00\01"}, %arg1: tensor<10x1024xf32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg2: tensor<128x1024xf32> {mhlo.sharding = ""}) {
    func.return
  }
}

include "mhlo/IR/hlo_ops.td"


// See the function doc in the header file.
def GetPadOpType : NativeCodeCall<
  "GetPadOpType((*$0.begin()).getDefiningOp<mhlo::CompositeOp>())">;

// See the function doc in the header file.
def GetAvgPoolOpPadAttr: 
  NativeCodeCall<"GetAvgPoolOpPadAttr($_builder, (*$0.begin()).getDefiningOp<mhlo::CompositeOp>())">;

      return %1: tensor<2048xi32>
   }
   // CHECK-LABEL: func @tpu_func
   // CHECK: mhlo.type_extensions
   func.func @tpu_func (
      %arg0: tensor<2048xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg1: tensor<2048xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) -> (tensor<2048xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) {
  // TODO(b/292540052): Below tf.addV2 instruction is replaced with just