See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Declare bridge passes that perform legalization and/or lowering.

include "mlir/Pass/PassBase.td"

def ConvertTFQuantOpsToMHLO : Pass<"quant-convert-tf-quant-ops-to-mhlo", "mlir::func::FuncOp"> {
  let summary = "Convert TF Quant ops to MHLO quantizated ops.";

              (UpdateShapeWithAxis<-1> $qtype, $old_value))),
          [(CanUpdateShapeWithAxis<-1> $qtype, $old_value)]>;

// The axis is set to 0 because the transpose is from the legalization of
// tf.conv2d and the new channel axis is the first dimension.
def ReorderTransposeDequantQuantUsedByConv :
      Pat<(TF_TransposeOp:$old_value
              (TFL_DequantizeOp (TFL_QuantizeOp $input, $qtype)), $perm),

  if (options_.legalize_to_tflite_ops) {
    // After we creat the alternative subgraph, we can still do canonicalization
    // legalization & other optimizations as long as we're not inlining the
    // function.
    // And in fact, we probably need to do the proper legalization, for the
    // compute cost to work. (in case we added some TF ops)
    pass_manager->addPass(mlir::TFL::CreatePrepareTFPass(

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// This is the legalization pattern definition file for HLO to TFL.

include "mlir/IR/OpBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mhlo/IR/hlo_ops.td"
include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"

      TypeID::get<TF::OutfeedEnqueueTupleOp>(),
      TypeID::get<TF::XlaShardingOp>(),

      // These ops have undetermined bugs, may not be legalizable with
      // XlaOpKernel
      // legalization in TF2XLA fallback. By legalization with MLIR, we can fix
      // the bug. b/195583695 describes the motivation of this change.
      // See b/216355804 how to reproduce the bug regarding tf.RandomUniform Op

      // Conditional ops

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// TFLite variables legalization patterns
include "mlir/IR/PatternBase.td"
include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

def HasSupportedElementType :

  // TF -> StableHLO legalization.
  // Skip StatefulPartitionedCall to preserve aliased functions.
  mlir::odml::AddLegalizeTFToStablehloPasses(pm, /*skip_quantization_ops=*/true,
                                             /*skip_resize=*/false,
                                             /*skip_partitioned_calls=*/true);
  // StableHLO -> MHLO legalization for MHLO optimization.

      : quant_specs(std::move(specs)) {}

  // If `emit_builtin_tflite_ops` is true, TF Lite legalization passes will be
  // added, which produces TF Lite ops.
  bool emit_builtin_tflite_ops = true;
  // If `lower_tensor_list_ops` is true, tensorlist ops will be lowered to basic
  // TF ops before legalization to TF Lite dialect.
  bool lower_tensor_list_ops = false;

};

TEST_F(TestPassInstrumentation, CreatedCalledAndSetsPassName) {
  RegisterPassInstrumentor(kTestInstrumentationName, [&]() {
    return std::make_unique<TestInstrumentor>(this);
  });
  constexpr char legalization[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

//===----------------------------------------------------------------------===//
// tf.UniformQuantizedDotHybrid legalization
//===----------------------------------------------------------------------===//

// CHECK-LABEL: func @quantized_matmul_fn
func.func @quantized_matmul_fn(%input: tensor<?xf32>) -> tensor<?xf32> {