// RUN: tf-opt "-tfxla-verify-legalization=legalize-chlo=true" -verify-diagnostics -split-input-file %s | FileCheck -dump-input=fail %s
// Tests the VerifyTFXLALegalization Pass, that just ensures we don't have
// any illegal ops at the end of the pipeline.

// CHECK-LABEL: allowsMHLO
func.func @allowsMHLO() -> (tensor<8x64x32x4xcomplex<f32>> {mhlo.sharding = ""}) {
  %0 = mhlo.constant dense<(1.000000e+00,-1.000000e+00)> : tensor<128x32x4xcomplex<f32>>

// RUN: stablehlo-quant-opt %s -verify-quant-legalization -split-input-file -verify-diagnostics | FileCheck %s

// CHECK-LABEL: func @legal_i8
func.func @legal_i8(%arg0: tensor<1xi8>) -> tensor<1xi8> {
  func.return %arg0: tensor<1xi8>
}

// -----

func.func @illegal_qint8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint8> {
  // expected-error@+1 {{'func.return' op is illegal as it is a UQ op or contains uq/qint types}}

Tres Popp <******@****.***> 1648061556 -0700

// RUN: tf-opt "-tfxla-verify-legalization=legalize-chlo=false" -verify-diagnostics -split-input-file %s | FileCheck %s --dump-input=fail
// Tests the VerifyTFXLALegalization Pass, that just ensures we don't have
// any illegal ops at the end of the pipeline. This runs with
// legalize-chlo=false since errors can't be mixed with the legalize-chlo=True
// version.

// CHECK-LABEL: allows_chlo

// RUN: tf-opt -tfl-prepare-tf=allow-bf16-and-f16-type-legalization=true %s | FileCheck %s

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

// CHECK-LABEL: conv_2d_bf16
func.func @conv_2d_bf16(%arg0 : tensor<256x32x32x3xbf16>, %arg1 : tensor<3x3x3x16xbf16>) -> tensor<256x8x7x16xbf16> {

    Option<"prefer_tf2xla_", "prefer-tf2xla", "bool",
        /*default=*/"false",
        "Prioritize tf2xla fallback legalization over MLIR legalization "
        "patterns">,
  ];

  let constructor = "mlir::mhlo::createLegalizeTFPass()";
  let dependentDialects = [
    "arith::ArithDialect",
    "chlo::ChloDialect",
    "func::FuncDialect",
    "mhlo::MhloDialect",
    "quant::QuantizationDialect",
    "shape::ShapeDialect",

               "legalzation, with a preference toward TF2XLA.";
  } else if (tf2xla_fallback_device_type) {
    VLOG(1) << "TF to XLA legalization patterns include all native patterns "
               "and TF2XLA fallback patterns.";
  } else {
    VLOG(1) << "TF to XLA legalization patterns are native patterns only.";
  }

  // Set patterns to legalize_lower_patterns to check whether they should use

/// list.
void PopulateLegalizeTfPatterns(MLIRContext* context,
                                RewritePatternSet* patterns);

// Populates TF to MHLO legalization for some of the quantization ops.
//
// TODO(hinsu): Remove this once we combine quantized and non quantized op
// legalization in the ODML conversion pipeline.
void PopulateLegalizeTfQuantizationPatterns(MLIRContext* context,

// - TensorFlow passes
// - Legalization passes
// - MHLO passes
//
// device_type: XLA JIT device to use for compilation such as "XLA_CPU_JIT",
//   "XLA_GPU_JIT" or "XLA_TPU_JIT".
// enable_op_fallback: when this is true, prefer tf2xla fallback kernels over
// MLIR
//   native kernels for legalization to HLO.
// custom_legalization_passes: passes to run before the default TF legalization
//   passes for backend-specific ops.

  // this happen.

  // StableHLO -> MHLO legalization.
  pm.addPass(mhlo::createStablehloLegalizeToHloPass());

  AddMhloOptimizationPasses(pm, /*enable_stablehlo_quantizer=*/false);
  // TODO: b/293149194 - Add `createFoldBroadcastPass` back to
  // `AddMhloOptimizationPasses`
  pm.addNestedPass<func::FuncOp>(createFoldBroadcastPass());

  // MHLO -> StableHLO legalization.
  pm.addPass(mhlo::createHloLegalizeToStablehloPass());