limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_quantize_op.h"

#include <functional>
#include <optional>

#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/types/optional.h"

A. Unique TensorFlower <******@****.***> 1711265865 -0700

        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "tf_quantize_op",
    srcs = [
        "tf_quantize_op.cc",
    ],
    hdrs = ["tf_quantize_op.h"],
    compatible_with = get_compatible_with_portable(),
    deps = [
        "//tensorflow/compiler/mlir/quantization/common/quantization_lib",

limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_quantize_op.h"

#include <optional>

#include <gtest/gtest.h>
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/Dialect/Quant/QuantOps.h"  // from @llvm-project

#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_quantize_op.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

// "qtype" attribute, we should update the type shape in this attribute.
def ReorderTransposeDequantQuant :
      Pat<(TF_TransposeOp:$old_value
              (TFL_DequantizeOp (TFL_QuantizeOp $input, $qtype)), $perm),
          (TFL_DequantizeOp (TFL_QuantizeOp
                                (TF_TransposeOp $input, $perm),
                                (UpdateShapeWithAxis<-1> $qtype, $old_value))),

// squashed to a requantize op if the scales are different.
def : Pat<(TFL_QuantizeOp (TFL_DequantizeOp $in), $qt), (replaceWithValue $in)>;

// If the tfl.dequantize op wasn't fused, we shouldn't quantize the floating
// point constant.
def : Pat<(TFL_DequantizeOp
             (TFL_QuantizeOp (Arith_ConstantOp F32ElementsAttr:$cst), $qt)),
          (TFL_ConstOp $cst)>;

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

// Both Quantize and Dequantize ops have side effects, so we have to define
// patterns to remove dead ones after the quantization rewrite.
def : Pat<(TFL_QuantizeOp:$op $in, $qt), (replaceWithValue $in), [(HasNoUseOf:$op)]>;

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

                             (TFL_ExpandDimsOp $input, $dim)>;

def LegalizeFakeQuantToDequantizeQuantize : Pat<
  (TF_FakeQuantWithMinMaxArgsOp $inputs, $min, $max, $num_bits, $narrow_range),
  (TFL_DequantizeOp
   (TFL_QuantizeOp $inputs,
                   (ConvertToQuantTypeFromAttrs $inputs, $min, $max,
                    $num_bits, $narrow_range)))>;

def LegalizeFill : Pat<(TF_FillOp $arg, $value), (TFL_FillOp $arg, $value)>;