include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"

// Quantize attribute $0 by using quantization parameter from %1.
def QuantizeByQuantizedType : NativeCodeCall<"quant::Quantize($0, $1.getValue())">;
def F32ElementsAttr : ElementsAttrBase<
  CPred<"$_self.cast<ElementsAttr>().getShapedType().getElementType().isF32()">, "float constant tensor">;

// Squash tfl.dequantize and tfl.quantize pairs.

// After applying the function, a quantize/dequantize functions are created
// where the body of each function contains a specific quantization algorithm.
// The input of the quantize function has one operand of
// IsValueWithQuantizablePrecision and the output is a tensor with supported
// quantized precision (like int8). For dequantize function, it is the other way
// around.

    auto op_before_dequantize = original_dequantize_op.getOperand(0);

    // Create a new dequantize op that is propagated.
    rewriter.setInsertionPointAfter(user_op);
    TF::PartitionedCallOp new_dequantize_op =
        cast<TF::PartitionedCallOp>(rewriter.clone(*original_dequantize_op));

    // Skip the original dequant op and connect the op before dequantize to the
    // user op.
    user_op->setOperand(user_idx, op_before_dequantize);

            op->user_begin()->hasTrait<OpTrait::IsTerminator>())
          return failure();
      }
      // If the quantize op is a requantize op, it is being used in other scale
      // adjustments and should be kept. Instead, moving dequantize op before
      // the requantize op to remove the unnecessary requantize op.
      if (auto qtype =
              QuantizedType::getQuantizedElementType(q.getArg().getType())) {

  // CHECK:      {
  // CHECK-NEXT:  version: 3,
  // CHECK-NEXT:  operator_codes: [ {
  // CHECK-NEXT:    deprecated_builtin_code: 6,
  // CHECK-NEXT:    version: 1,
  // CHECK-NEXT:    builtin_code: DEQUANTIZE
  // CHECK-NEXT:  }, {
  // CHECK-NEXT:    deprecated_builtin_code: 4,
  // CHECK-NEXT:    version: 2,
  // CHECK-NEXT:    builtin_code: DEPTHWISE_CONV_2D
  // CHECK-NEXT:  } ],
  // CHECK-NEXT:  subgraphs: [ {

    // Finally, use the quantization parameter to create the quantize and
    // dequantize ops, and insert them between the tf.FakeQuantWithMinMaxVarsOp
    // and its users.
    auto quantize = rewriter.create<quantfork::QuantizeCastOp>(
        tf_op.getLoc(), qtype.getValue(), input);
    auto dequantize = rewriter.create<quantfork::DequantizeCastOp>(
        tf_op.getLoc(), res_type, quantize.getResult());

  // CHECK:      {
  // CHECK-NEXT:  version: 3,
  // CHECK-NEXT:  operator_codes: [ {
  // CHECK-NEXT:    deprecated_builtin_code: 6,
  // CHECK-NEXT:    version: 1
  // CHECK-NEXT:    builtin_code: DEQUANTIZE
  // CHECK-NEXT:  }, {
  // CHECK-NEXT:    deprecated_builtin_code: 4,
  // CHECK-NEXT:    version: 1
  // CHECK-NEXT:    builtin_code: DEPTHWISE_CONV_2D
  // CHECK-NEXT:  } ],
  // CHECK-NEXT:  subgraphs: [ {

      if (!q->getAttr(kVolatileOpAttrName)) return failure();

      // If the quantize op is a requantize op, it is being used in other scale
      // adjustments and should be kept. Instead, move dequantize op before the
      // requantize op to remove the unnecessary requantize op.
      if (const QuantizedType qtype =
              QuantizedType::getQuantizedElementType(q.getArg().getType())) {

// and tfl.dequantize pairs before tf.FakeQuant* being foled.
LogicalResult ConvertFakeQuantOps(func::FuncOp func, MLIRContext* ctx,
                                  bool use_fake_quant_num_bits) {
  OpBuilder builder(func);
  if (failed(UnwrapTFCustomOps(func, builder))) {
    return failure();
  }

  // Insert the tfl.quantize/tfl.dequantize ops after the tf.FakeQuant* ops to

                  (Arith_ConstantOp ConstantAttr<RankedF32ElementsAttr<[]>,
                   "-1.0f">), TFL_AF_None), $act)>;

// Squash tfl.dequantize and tfl.quantize pairs.
// TODO(b/185915462): Compare the scale of input and output. This can also be
// squashed to a requantize op if the scales are different.