// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-unwrap-xla-call-module-op | FileCheck %s

// Tests if XlaCallModule op without quantizable trait that calls function with
// '_from_xla_call_module' trait is unwrapped.
// Tests if XlaCallModule op with quantizable trait is not unwrapped.
// Tests if XlaCallModule op without quantizable trait that calls function
// without '_from_xla_call_module' trait is not unwrapped.

constexpr llvm::StringRef kStablehloMainFunctionName = "main";

// After deserializing the stablehlo functions from XlaCallModule,
// this XlaCallModule attribute refers to the deserialized stablehlo main
// function.
constexpr llvm::StringRef kStablehloEntryFunctionAttrName = "_entry_function";

// Every stablehlo function deserialized from XlaCallModule has this attribute.
constexpr llvm::StringRef kFromXlaCallModuleAttrName = "_from_xla_call_module";

  }
}

// CHECK-LABEL: @conv_3_fn
// CHECK: tf.XlaCallModule
// CHECK-SAME: _entry_function = @composite_conv_fn_6, _original_entry_function = "composite_conv_fn_6"
// CHECK: tf.XlaCallModule
// CHECK-SAME: _entry_function = @composite_conv_fn_5, _original_entry_function = "composite_conv_fn_5"

// CHECK-LABEL: @conv_1_fn
// CHECK: tf.XlaCallModule

// -----

// Tests composite tf.XlaCallModule is converted to func.call.

module {
  // CHECK-LABEL: func.func @main
  func.func @main(%arg0: tensor<1x1024xf32>) -> tensor<1x3xf32> {
    // CHECK: call @composite_dot_general_fn_1
    // CHECK-SAME: (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
    // CHECK-NOT: tf.XlaCallModule

  // CHECK-SAME: %[[ARG0:[^:[:space:]]+]]
  // CHECK-SAME: %[[ARG1:[^:[:space:]]+]]
  func.func private @serving_default(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>) -> tensor<1x3xf32> {

// between constant and XlaCallModule op with empty `weight_only_ptq` method
// and function name containing conv.

func.func @qdq_for_conv_weight_empty(%arg0: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32> attributes {tf._original_func_name = "main_0"} {
  %cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
  %0 = "tf.XlaCallModule"(%arg0, %cst) {

    %2 = "tf.XlaCallModule"(%arg0, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_dot_general_with_bias_and_relu6_dynamic_fn_2, _original_entry_function = "composite_dot_general_with_bias_and_relu6_dynamic_fn_2", _tfl_quant_trait = "fully_quantizable"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>

  // CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_1]], %[[SUBGRAPH_1]]) <{Sout = [#tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1"
  // CHECK: %[[CUSTOM_AGGREGATOR_2:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])

// RUN: tf-opt %s -split-input-file -tf-xla-call-module-deserialization | FileCheck %s

// Tests that `tf.XlaCallModule` with both StableHLO module and TF function
// calls can be deserialized.

// CHECK-LABEL: module
module {
  // CHECK-LABEL: func private @_tf_func
  func.func private @_tf_func(%arg0: tensor<?xi32>, %arg1: tensor<*xi32>) {
    // CHECK: tf.StreamResults

    // StreamResults is a pseudo op in this test.

}

// -----

// Tests that the output of the tf.XlaCallModule op has been replaced by
// a quantized type, and the corresponding quantfork.qcast ops that turned
// the float output to a quantized type is removed.

// CHECK-LABEL: quantize_simple_xla_call_module_no_operand
func.func private @quantize_simple_xla_call_module_no_operand() -> tensor<1x3xf32> {