// RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions -quant-quantize-composite-functions | FileCheck %s

module {
  func.func @conv(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>, tensor<*xf32>) {

    StringRef function_name = entry_function.getValue();
    if (!function_name.starts_with("composite_")) {
      return spec;
    }

    if (function_name.contains("conv")) {
      // Looks up `Method` to see if it should be per-channel quantized and
      // populates the spec accordingly.
      PopulateCoeffOpQuantDimIfPerChannelQuantized(call_op, *spec);

        # One pure conv
        out = nn_ops.conv2d(
            out,
            self.conv_filters,
            strides=(1, 1, 2, 1),
            dilations=(1, 1, 1, 1),
            padding='SAME',
            data_format='NHWC',
        )

        # One fakequant attached conv
        if is_qat_model:
          out = array_ops.fake_quant_with_min_max_args(

	n = DefaultLit(n, types.Types[types.TINT])
	*np = n

	return true
}

func Conv(n ir.Node, t *types.Type) ir.Node {
	if types.IdenticalStrict(n.Type(), t) {
		return n
	}
	n = ir.NewConvExpr(base.Pos, ir.OCONV, nil, n)
	n.SetType(t)
	n = Expr(n)
	return n
}

// ConvNop converts node n to type t using the OCONVNOP op

	OCALLMETH:         8,
	OCALL:             8,
	OCAP:              8,
	OCLEAR:            8,
	OCLOSE:            8,
	OCOMPLIT:          8,
	OCONVIFACE:        8,
	OCONVNOP:          8,
	OCONV:             8,
	OCOPY:             8,
	ODELETE:           8,
	OGETG:             8,
	OLEN:              8,
	OLITERAL:          8,
	OMAKESLICE:        8,
	OMAKESLICECOPY:    8,
	OMAKE:             8,

    const auto module_op = op->getParentOfType<ModuleOp>();
    const SymbolTable symbol_table(module_op);
    func::FuncOp func = symbol_table.lookup<func::FuncOp>(function_name);

    if (function_name.contains("conv")) {
      return (*(func.getOps<mlir::stablehlo::ConvolutionOp>().begin()))
          .getDimensionNumbers()
          .getKernelOutputFeatureDimension();
    } else if (function_name.contains("dot_general")) {

  ASSERT_THAT(new_config.specs().specs(), SizeIs(1));

  const QuantizationSpec& spec = new_config.specs().specs(0);
  EXPECT_THAT(spec.matcher().function_name().regex(),
              StrEq("^.*(conv|dot_general).*"));
  EXPECT_TRUE(spec.method().has_weight_only_ptq());

  const WeightOnlyPtq& weight_only_ptq_spec = spec.method().weight_only_ptq();

  EXPECT_THAT(weight_only_ptq_spec.input_quantized_types(),

// or `std::nullopt` if the given op is not per-channel quantizable.
std::optional<int64_t> GetDotGeneralQuantizationDim(
    ::mlir::stablehlo::DotGeneralOp dot_general_op);

// Checks if a `StringRef` contains 'conv' or 'dot_general'.
bool ContainsConvOrDot(StringRef str);

}  // namespace mlir::quant

}

func mustGetLabels(t test.Failer, obj object.K8sObject, path string) map[string]string {
	t.Helper()
	got := mustGetPath(t, obj, path)
	conv, ok := got.(map[string]any)
	if !ok {
		t.Fatalf("could not convert %v", got)
	}
	ret := map[string]string{}
	for k, v := range conv {
		sv, ok := v.(string)
		if !ok {
			t.Fatalf("could not convert to string %v", v)
		}
		ret[k] = sv
	}
	return ret
}

// CHECK-NEXT:  %[[conv:.*]] = "tfl.conv_2d"(%arg0, %[[q_cst_0]], %[[q_cst_1]]) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}>
// CHECK-NEXT:  %[[reshape:.*]] = "tfl.reshape"(%[[conv]], %[[cst]]) : (tensor<1x112x112x32x!quant.uniform<u8:f32, 0.023528476789885875>>, tensor<2xi32>)