# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Generates the quantized function library contained header file."""

import ast
import re
import string
from typing import Sequence

from absl import app
from absl import flags

// checkpoint saving and restoring. This function returns a `SaverDef` instance
// with four fields populated: `version`, `filename_tensor_name`,
// `restore_op_name` and `save_tensor_name`. For valid quantized `graph_def` and
// `control_ret_node_names`, it should be able to retrieve the last three fields
// if there is at lest one variable in the graph.
//

#include "mlir/IR/Builders.h"  // from @llvm-project

namespace mlir::quant {

// Caclulate padding values for XLA ops.
// Padding values for Uniform Quantized ops can be generated with this method as
// well as it shares the same definition for padding attribute with the XLA ops.
Value CalculatePaddingAndPadIfNeeded(OpBuilder &builder, Location loc,

    func.return %dot_out : tensor<*x!tf_type.qint32>
  }

  // Quantize initial input at the start of the graph. Output is qint8.
  func.func @quantize_i8(%input : tensor<*xf32>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>) -> tensor<*x!tf_type.qint8> {
    %quantize = "tf.UniformQuantize"(%input, %input_scale, %input_zp) {
      Tin = "tfdtype$DT_FLOAT",
      Tout = "tfdtype$DT_QINT8",

          clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
          clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
                     "Uses TF Uniform Quantized ops"))};

  Option<int64_t> min_num_elements_for_weights_{
      *this, "min-num-elements-for-weights", llvm::cl::init(0),
      llvm::cl::desc("The minimum required number of elements in a weight "

  auto func_op = module_op->lookupSymbol<func::FuncOp>("quantize");
  ASSERT_THAT(func_op, NotNull());

  auto uniform_quantize_op_itr =
      func_op.getBody().op_begin<mlir::stablehlo::UniformQuantizeOp>();
  ASSERT_THAT(
      uniform_quantize_op_itr,
      Ne(func_op.getBody().op_end<mlir::stablehlo::UniformQuantizeOp>()));

  // `uniform_quantize` is considered partially quantized because its output is

// RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions='quantization-method=weight_only target-opset=XLA' -quant-quantize-composite-functions='quantization-method=weight_only target-opset=XLA' -symbol-dce | FileCheck --check-prefix=PerTensor %s

      Constraint<CPred<"IsEinsumSupportedByXlaDotV2($0)">>;

// This attribute can be used in the `AttributeList` for missing attributes. It
// is necessary to keep other attributes in the same index as the quantized
// composite function.

                            "quant::QuantizationDialect"];

}

def OptimizeIntGraph : Pass<"optimize-int-graph", "mlir::func::FuncOp"> {
  let summary = "Optimization patterns for quantized integer graph";

  let description = [{
    This includes patterns for merging addition of zp offset and bias.
  }];

  let constructor = "mlir::quant::stablehlo::CreateOptimizeIntGraphPass()";

  const double qmaxDouble = qmax;
  scale = (rmax - rmin) / (qmaxDouble - qminDouble);

  // Zero point computation.
  // In float, solve the affine equation for any known pair
  // (real value, corresponding quantized value), of which, two such pairs
  // are known: (rmin, qmin), (rmax, qmax).
  // The arithmetic error on the zero point computed from either pair will be
  // roughly machine_epsilon * (sum of absolute values of terms).