// RUN: tac-opt-all-backends -tfl-pick-subgraphs %s -split-input-file -verify-diagnostics | FileCheck %s

module {
  func.func @main(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>, %arg2: tensor<100xf32>, %arg3: tensor<100xf32>) -> tensor<2x100xf32> {
    %0 = func.call @func_0_GPU_FLOAT(%arg0, %arg1, %arg2) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<100xf32>, tensor<100xf32>, tensor<100xf32>) -> tensor<100xf32>

// RUN: tac-opt-all-backends -tfl-raise-target-subgraphs %s -split-input-file | FileCheck %s
// RUN: tac-opt-all-backends -tfl-raise-target-subgraphs="skip-raise-cpu-ops=true" %s -split-input-file | FileCheck %s --check-prefixes=CHECK-SKIP-CPU
// RUN: tac-opt-all-backends -tfl-raise-target-subgraphs="ignore-inference-type=true" %s -split-input-file | FileCheck %s --check-prefixes=CHECK-IGNORE-INFERENCE-TYPE

module {
func.func @simpleWhile(%arg0: tensor<i32>) -> tensor<i32> {

//     subgraph3 (GPU)      subgraph4(GPU)
//         |             /
//      subgraph5 (GPU)
//         |
//      subgraph6 (CPU)
//
// The overall workflow of the pick subgraphs pass:
//  1) Build subgraphs
//    1.1) Collect output subgraphs.
//    1.2) Build `Subgraph` and their "alternative view" from FuncOp.
//  2) Pick subgraphs
//    2.1) Populate the "dp table" for (subgraph, hardware).

  ASSERT_THAT(status, Eq(kTfLiteOk));

  const auto subgraph = model_.subgraphs[0].get();

  // The model should only have 3 inputs and 1 output.
  EXPECT_THAT(subgraph->inputs, SizeIs(3));
  EXPECT_THAT(subgraph->outputs, SizeIs(1));

  const auto& op1 = subgraph->operators[1].get();
  const auto& op2 = subgraph->operators[2].get();
  const auto& op3 = subgraph->operators[3].get();

    for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
         ++subgraph_idx) {
      const auto subgraph = model->subgraphs()->Get(subgraph_idx);
      for (size_t i = 0; i < subgraph->inputs()->size(); ++i) {
        if (subgraph->inputs()->Get(i) == tensor_idx) {
          return true;
        }
      }
      for (size_t i = 0; i < subgraph->outputs()->size(); ++i) {

        GetInferenceDeviceTypeForOp(partition_ops.front()).value();
    Subgraph old_subgraph(partition_ops, ++func_count);
    OpsAdded ops_added;
    ExtractSubgraphToFunc(old_subgraph, builder, module, ops_added);
    AddAttrs(ops_added, builder, func_count);
    // Ops in "CPU" subgraphs may nested regions with other device subgraphs.
    // We recur into these nested blocks to raise those as well. We don't raise

so we will annotate the op with "CPU" since "GPU" does not support this op.

#### Raise Target Subgraphs Pass

In this pass, ops will be broken down into subgraph. Those ops have the same
target annotation will be raised as subgraphs.

In this pass, subgraph is actually implemented with `FuncOp`.

Take the following code as an example:

```

          fb_builder.CreateVector(std::vector<float>({-1.0, 2.0}))),
  };
  const auto subgraphs = {CreateSubgraphMetadata(
      fb_builder, fb_builder.CreateVector(ops.begin(), ops.size()))};

  const auto metadata = CreateRuntimeMetadata(
      fb_builder, hardwares,
      fb_builder.CreateVector(subgraphs.begin(), subgraphs.size()));
  fb_builder.Finish(metadata);

  return std::string(

}

// Check if the op should be added to the subgraph.
// The op should be added to the subgraph if all of its users match one
// of following two conditions:
// 1: The user is already in the current subgraph.
// 2: The user will reach a dead end.
//
// If the op should be added to the subgraph and there are users who
// will reach the dead end, add the ops on the dead end to the subgraph as well.
bool ShouldAddOpToSubgraph(Operation* op,

// CHECK: #[[LOC2:.*]] = loc("output1"(#[[LOC0]]))
// CHECK: #[[SPLIT_LOC]] = loc(fused[#[[LOC1]], #[[LOC2]]])

{
  "version": 3,
  "operator_codes": [
    {
      "builtin_code": "SPLIT"
    }
  ],
  "subgraphs": [
    {
      "tensors": [
        {
          "shape": [
            1
          ],
          "name": "split_dim",
          "quantization": {
          }
        },
        {