// 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 enable-per-channel-quantization=true' -symbol-dce...

// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op='target-opset=XLA quantization-method=weight_only enable-per-channel-quantization=false' | FileCheck --check-prefix PerTensor %s
// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op='target-opset=XLA quantization-method=weight_only enable-per-channel-quantization=true' | FileCheck --check-prefix PerChannel %s

module {
  // For XLA weight-only per-channel depthwise convolution, tensor shape should have

// RUN: tf-quant-opt %s -quant-insert-quantized-functions='quantization-method=weight_only target-opset=XLA' | FileCheck %s

// Empty module
module {
  func.func @simple_fn(%arg0: tensor<*xf32>) -> tensor<*xf32> {
    func.return %arg0 : tensor<*xf32>
  }
}

// CHECK-NOT: func private @internal_dequantize_f32
// CHECK-NOT: func private @internal_conv3d_fn
// CHECK-NOT: func private @internal_batch_matmul_fn

          clEnumValN(tensorflow::quantization::QuantizationMethod::
                         METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8,
                     "weight_only", "Post-training weight_only quantizaiton"))};
};

class CheckQuantizableOps
    : public mlir::OpRewritePattern<TF::PartitionedCallOp> {
 public:
  explicit CheckQuantizableOps(MLIRContext* context,

  std::string dynamic_quant_kernel_support_regex =
      "bool GetDynamicRangeQuantKernelSupport() { return true; }";
  raw_ostream &os = *ostream;
  std::vector<std::string> weight_only;
  llvm::sort(defs, LessRecord());

  os.indent(0) << "const std::set<std::string> &ExportDynamicRangeSpec() {\n";
  os.indent(2) << "static const std::set<std::string> * result =\n";

          clEnumValN(tensorflow::quantization::QuantizationMethod::
                         METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8,
                     "weight_only", "Post-training weight_only quantizaiton"))};

  Option<OpSet> op_set_{
      *this, "target-opset", llvm::cl::init(OpSet::TF),
      llvm::cl::desc("Choose target opset."),
      llvm::cl::values(

      break;
    // Note: This is weight-only quantization by default, but with the legacy
    // flag "--force_dynamic_range_in_kernel", a DRQ behavior will be forced
    // in the kernel.
    case PresetQuantizationMethod::WEIGHT_ONLY:
      weight_component = custom_method.add_quantization_component_spec();
      SetQuantizationComponentSpec(weight_component,
                                   QuantizationComponentSpec::COMPONENT_WEIGHT,

    // Apply default weight-only quantization. Weights are quantized during
    // conversion, then dequantized during inference.
    // Activation: f32, Weight: qi8, Bias: f32
    WEIGHT_ONLY = 1;

    // Apply default dynamic range quantization. Quantized tensor value's
    // ranges are determined during graph runtime.
    // Activation: f32, Weight: qi8, Bias: f32

// CHECK-LABEL: private @composite_conv3d_fn_1

// WEIGHTONLY-DAG: %[[CST:.*]] = "tf.Const"() {{.*}} : () -> tensor<2x3x3x3x2xf32>
// WEIGHTONLY: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]])
// WEIGHTONLY-SAME: f = @composite_conv3d_fn_1}>
// WEIGHTONLY: {_tfl_quant_trait = "fully_quantizable"
// WEIGHTONLY: %[[RELU:.*]] = "tf.Relu"(%[[PARTITIONEDCALL_0]])
// WEIGHTONLY: return %[[RELU]]

                     "drq", "Post-training dynamic-range quantizaiton"),
          clEnumValN(tensorflow::quantization::QuantizationMethod::
                         METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8,
                     "weight_only", "Post-training weight-only quantizaiton"))};

  Option<bool> enable_per_channel_quantization_{
      *this, "enable-per-channel-quantization", llvm::cl::init(false),