// RUN: tf-quant-opt %s -mark-functions-noinline='noinline-functions=noinline0' \
// RUN:     -allow-unregistered-dialect -mlir-disable-threading \
// RUN:     -split-input-file -verify-diagnostics | FileCheck %s

// Tests that the function is marked tf._noinline = true.

// CHECK-LABEL: @noinline0
// CHECK-SAME: attributes {{{.*tf._noinline = true.*}}}
func.func @noinline0() -> (tensor<0xf32>) {

// Required when using LLVM_DEBUG macro.
#define DEBUG_TYPE "mark-functions-noinline"

namespace mlir {
namespace quant {
namespace {

// Name of the boolean attribute indicating whether the function can be
// inlined or not.
constexpr StringRef kTfNoinlineAttr = "tf._noinline";

// This pass marks functions with the attribute `tf._noinline = true` so that

// RUN: tf-opt %s -tf-strip-noinline-attribute | FileCheck %s

// CHECK-LABEL: func @strip_simple(
// CHECK-NOT: tf._noinline
func.func @strip_simple() -> tensor<2xi32> attributes {tf._noinline = true} {
  // CHECK-NEXT: %[[CST:.*]] = "tf.Const"
  %cst = "tf.Const"() { value = dense<2> : tensor<2xi32> } : () -> tensor<2xi32>
  func.return %cst : tensor<2xi32>

// This tranformation pass strips any "_noinline" attributes from the module.
struct StripNoinlineAttributePass
    : public impl::StripNoinlineAttributePassBase<StripNoinlineAttributePass> {
 public:
  // void runOnOperation() override;
  void runOnOperation() override {
    // Strip the "tf._noinline" attribute from top-level functions.
    for (auto func_op : getOperation().getOps<func::FuncOp>())

  func.func private @some_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> attributes {tf._noinline = true} {
    return %arg0 : tensor<1x2x2x3xf32>
  }

  // CHECK-SKIP: func.func private @some_other_func
  func.func private @some_other_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> attributes {tf._noinline = true} {
    return %arg0 : tensor<1x2x2x3xf32>
  }

// Check that we only outline the outermost partitioned call with
// `_xla_compile_device_type` in nested calls to a device cluster. The callee
// of any partitioned call outside of a device cluster is marked with
// `tf._noinline = true` to prevent it getting inlined to perseve runtime
// performance.
// CHECK-LABEL: func.func @nested_calls
func.func @nested_calls(%arg0: tensor<i32>) -> tensor<i32> attributes {tf.entry_function = {}} {

  func.return %1 : tensor<i32>
}

func.func private @xla_func_0(%arg0: tensor<1x3xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._XlaMustCompile = true, tf._noinline = true, tf._original_func_name = "should_not_be_used"} {
  %1 = "tf.AddV2"(%arg0, %arg1) : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
  %2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>

  %1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<f32>>>) -> tensor<f32>
  %2 = "tf.XlaLaunch"(%arg0, %1) {_noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<7xf32>, tensor<f32>) -> tensor<10xf32>
  func.return %2 : tensor<10xf32>

  %cst = "tf.Const"() { value = dense<2> : tensor<2xi32> } : () -> tensor<2xi32>
  func.return %cst : tensor<2xi32>
}

func.func private @simple_callee_with_noinline() -> tensor<2xi32> attributes {tf._noinline = true} {
  %cst = "tf.Const"() { value = dense<2> : tensor<2xi32> } : () -> tensor<2xi32>
  func.return %cst : tensor<2xi32>
}

// Test that simple TF operations can be inlined.

// CHECK-LABEL: func @inline_simple(

    return %arg0 : tensor<1x2x2x3xf32>
  }
}

// CHECK: module
// CHECK: tf.StatefulPartitionedCall
// CHECK: func.func private @some_func
// CHECK-NOT: func.call

// -----

// This test makes sure functions without tf._noinline=true is inlined.

module {