#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/hlo_matchers.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {

LogicalResult MatchReduceToArgMinMaxType1(mhlo::ReduceOp reduce_op,
                                          bool is_float, bool is_argmax);

LogicalResult MatchReduceToArgMinMaxType2(mhlo::ReduceOp reduce_op,
                                          bool is_argmax);

// Legalizes TF/XLA communication ops (TF dialect) to HLO dialect communication
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCommunicationPass();

// Legalizes TF/XLA collective ops (TF dialect) to HLO dialect collective
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCollectivePass();

// Verifies that the TF/XLA ops have all been lowered to MHLO.

#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {

class ConvertCustomCallOp : public OpConversionPattern<mhlo::CustomCallOp> {
 public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(
      mhlo::CustomCallOp mhlo_custom_call, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final;

#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "tensorflow/core/lib/monitoring/counter.h"
#include "tensorflow/core/platform/errors.h"

namespace mlir {
namespace mhlo {

namespace {

#define GEN_PASS_DEF_VERIFYTFXLALEGALIZATION
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"

#include "xla/layout.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/shape.h"
#include "xla/stream_executor/tpu/c_api_conversions.h"
#include "xla/stream_executor/tpu/tpu_api.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"

namespace mlir {
namespace mhlo {
namespace {

#define GEN_PASS_DEF_INFEEDOPSXLAADJUSTLAYOUT

    func.func @main(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>> {
      %0 = "tf.Identity"(%arg0) : (tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
      func.return %0 : tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
    }
  })";
  SetPassThatChangedIdentity("");

#include "xla/mlir_hlo/mhlo/transforms/passes.h"

namespace mlir::quant::stablehlo {

void AddQuantizationLoweringPasses(mlir::OpPassManager& pm) {
  // These passes are grouped together and must run in this specific order.
  pm.addNestedPass<mlir::func::FuncOp>(CreateConvertTFQuantOpsToMHLOPass());
  pm.addNestedPass<mlir::func::FuncOp>(mhlo::createChloLegalizeToHloPass());

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

}  // namespace mhlo
}  // namespace mlir

// RUN: tf-opt -tf-to-hlo-pipeline %s | FileCheck %s

// Verifies that constants generated post shape inference are propagated.
// get_shape result in this test.
module attributes {tf.versions = {producer = 179 : i32}} {

  // CHECK-LABEL: func @main
  func.func @main(%arg0: tensor<10x19xf32>, %arg1: tensor<19x10xf32> {mhlo.is_same_data_across_replicas = true}) -> tensor<?xi64> {
    %0 = "tf.Shape"(%arg0) : (tensor<10x19xf32>) -> tensor<2xi64>

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

def IsDefinedByConvOrDotGeneralOp : Constraint<Or<[
  CPred<"$0.getDefiningOp<mhlo::ConvolutionOp>()">,
  CPred<"$0.getDefiningOp<mhlo::DotGeneralOp>()">]>>;

def IsNull : Constraint<CPred<"!$0">>;

// This pattern optimizes: