// 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

// tensor to sparse format.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDenseToSparsePass();

// Creates function pass to legalize TF While to TFL While.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFWhilePass();

// Legalize tflite flex ops to TF ops.
std::unique_ptr<OperationPass<func::FuncOp>> CreateLiftTfliteFlexOpsPass();

  ];
}

def LegalizeWhilePass : Pass<"tfl-legalize-tf-while", "mlir::ModuleOp"> {
  let summary = "Legalize from TensorFlow While to TensorFlow Lite While.";
  let constructor = "CreateLegalizeTFWhilePass()";
  let dependentDialects = ["TFL::TensorFlowLiteDialect"];
}

def LegalizeVariablesPass : Pass<"tfl-legalize-variables-tf", "mlir::ModuleOp"> {
  let summary = "Legalize TensorFlow variables to TensorFlow Lite dialect.";

// RUN: tf-opt --tfl-legalize-tf-while %s -o - | flatbuffer_translate -mlir-to-tflite-flatbuffer - -o -  | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s

// CHECK-LABEL: main
func.func @main(%arg0: tensor<*xf32>) -> tensor<*xf32> {
  // CHECK:     %{{.*}} = tfl.add %{{.*}}, %{{.*}} {fused_activation_function = "NONE"} : tensor<*xf32>
  // CHECK:     return %{{.*}} : tensor<*xf32>

// RUN: tf-tfrt-opt -tf-to-tfrt %s -split-input-file -verify-diagnostics

// expected-error @+1 {{failed to legalize operation 'func.func' that was explicitly marked illegal}}
func.func @test_identity_wrong_type(%arg0: tensor<4x2x!tf_type.string>) -> tensor<4x2x!tf_type.stringref> {
  %0 = "tf.SomeOp"(%arg0) : (tensor<4x2x!tf_type.string>) -> tensor<4x2x!tf_type.stringref>
  func.return %0 : tensor<4x2x!tf_type.stringref>

        "prepare-tf.mlir": "medium",
        "prepare-tf-fake-quant.mlir": "medium",
        "prepare-tf-fake-quant-4bit.mlir": "medium",
        "raise-custom-ops.mlir": "medium",
    },
    tags_override = {
        "legalize-tf.mlir": ["no_rocm"],
        "optimize.mlir": ["no_rocm"],
        "prepare-tf.mlir": ["no_rocm"],
        "const-fold.mlir": ["no_mac_arm64"],
    },
    test_file_exts = ["mlir"],
)

    pass_manager->addPass(mlir::TF::CreateTFShapeInferencePass());
  }

  // Legalize while early to allow further constant folding.
  // TODO(jpienaar): This may not actually matter as we do canonicalization
  // after the legalize below, for now it needs to be below the above passes
  // that work on TF dialect and before inliner so that the function calls in

#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project

namespace mlir {
namespace TF {

// Legalize the tf ops to the quant ops, so the quantization passes can work.
std::unique_ptr<OperationPass<func::FuncOp>> CreateLegalizeTFToQuantPass();

// Fallbacks ops that are not supported by TF Quantization to TFLite Flex ops.

// RUN: tf-opt --tfl-legalize-tf-while %s -o - | flatbuffer_translate -mlir-to-tflite-flatbuffer - -o -  | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s

// This test is to test for unranked function output from input, the output type should be compatible with input type.

// CHECK: func @main(%arg0: tensor<1xf32>) -> tensor<*xf32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.while"}} {
// CHECK:   %0 = "tfl.while"(%arg0) ({

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