// RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions -quant-quantize-composite-functions='quantization-method=ptq target-opset=XLA' | FileCheck %s
// RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions -quant-quantize-composite-functions='quantization-method=ptq target-opset=XLA enable-per-channel-quantization=true' | FileCheck --check-prefix=PerChannel %s

module {

// CHECK-DAG: %[[w_quant:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}40, 20]
// CHECK-DAG: {{\[\[\[}}-87, -42]

// CHECK: %[[quantize:.*]] = "tf.PartitionedCall"(%arg0, %[[in_scale]], %[[in_zp]])
// CHECK-SAME: f = @quantize_i8
// CHECK: %[[conv_quant:.*]] = "tf.PartitionedCall"(%[[quantize]], %[[w_quant]], %[[b_quant]],
// CHECK-SAME: %[[in_scale]], %[[in_zp]], %[[w_scale]], %[[w_zp]],

	for i := 0; i < samples; i++ {
		window := time.Duration(math.Exp(float64(i)/(samples-1)*(logMax-logMin) + logMin))
		if quantiles == nil {
			plot[i] = make([]float64, 2)
			plot[i][1] = mmuCurve.MMU(window)
		} else {
			plot[i] = make([]float64, 1+len(quantiles))
			copy(plot[i][1:], mmuCurve.MUD(window, quantiles))
		}
		plot[i][0] = float64(window)
	}

	// Create JSON response.

    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).
    return "quant-quantize-weights";
  }

  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Quantize weights used by quantizable ops.";
  }

  void getDependentDialects(DialectRegistry& registry) const override {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"

//===----------------------------------------------------------------------===//
// The prepare-quantize-drq Pass.
//
namespace mlir {
namespace quant {

namespace {

using QuantizationUnit = std::pair<Operation*, int>;
using QuantizationUnits = llvm::SetVector<QuantizationUnit>;

  // Inserts the Quantize ops for requantizing a block argument.
  void RequantizeArg(BlockArgument arg, RequantizeStates& states);

  // Inserts the Quantize and Dequantize ops to quantize the value and returns
  // the Quantize op.
  void RequantizeValue(Value value, RequantizeStates& states, Location loc);

  // Returns the quantization parameter satisfies the same scale

func.func @modified(%arg0: tensor<1x224x224x3xf32>) -> tensor<1x401408xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "input", outputs = "output"}} {
  %cst = arith.constant dense<[1, 401408]> : tensor<2xi32>
  %0 = "tfl.quantize"(%arg0) {qtype = tensor<1x224x224x3x!quant.uniform<i8:f32, 7.812500e-03>>} : (tensor<1x224x224x3xf32>) -> tensor<1x224x224x3x!quant.uniform<i8:f32, 7.812500e-03>>

  let summary = "Restores function name from XlaCallModule op.";
}

def QuantizeCompositeFunctionsPass : Pass<"stablehlo-quantize-composite-functions", "ModuleOp"> {
  let summary = "Quantize composite functions with QDQ input / outputs.";
  let options = [
    Option<"enable_per_channel_quantized_weight_",
        "enable-per-channel-quantized-weight",

#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"

//===----------------------------------------------------------------------===//
// The post-quantize Passes.
//
namespace mlir {
namespace TFL {
namespace {
#define GEN_PASS_DEF_POSTQUANTIZEPASS
#define GEN_PASS_DEF_POSTQUANTIZEREMOVEQDQPASS
#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

// CHECK: {
// CHECK-NEXT:  version: 3,
// CHECK-NEXT:  operator_codes: [ {
// CHECK-NEXT:    deprecated_builtin_code: 114,
// CHECK-NEXT:    version: 1,
// CHECK-NEXT:    builtin_code: QUANTIZE
// CHECK-NEXT:  }, {
// CHECK-NEXT:    deprecated_builtin_code: 3,
// CHECK-NEXT:    version: 1,
// CHECK-NEXT:    builtin_code: CONV_2D
// CHECK-NEXT:  }, {
// CHECK-NEXT:    deprecated_builtin_code: 22,