text: /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
  
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  
      http://www.apache.org/licenses/LICENSE-2.0
  
  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
  ==============================================================================*/
  
  #include "tensorflow/lite/delegates/gpu/common/gpu_model.h"
  
  #include <algorithm>
  #include <utility>
  
  #include "absl/container/flat_hash_map.h"
  #include "absl/container/flat_hash_set.h"
  #include "tensorflow/lite/delegates/gpu/common/operations.h"
  #include "tensorflow/lite/delegates/gpu/common/selectors/operation_selector.h"
  #include "tensorflow/lite/delegates/gpu/common/selectors/special_selector.h"
  #include "tensorflow/lite/delegates/gpu/common/selectors/subgraph.h"
  #include "tensorflow/lite/delegates/gpu/common/task/serialization_base.h"
  #include "tensorflow/lite/delegates/gpu/common/task/storage_type_util.h"
  #include "tensorflow/lite/delegates/gpu/common/transformations/add_bias.h"
  #include "tensorflow/lite/delegates/gpu/common/transformations/global_pooling_to_reduce_op.h"
  #include "tensorflow/lite/delegates/gpu/common/transformations/merge_padding_with.h"
  
  #include "tensorflow/lite/tools/logging.h"
  
  namespace tflite {
  namespace gpu {
  
  namespace {
  bool IsReady(const absl::flat_hash_set<ValueId>& ready_tensors,
               const GpuNode& node) {
    for (const ValueId in_id : node.inputs) {
      if (ready_tensors.find(in_id) == ready_tensors.end()) {
        return false;
      }
    }
    return true;
  }
  
  absl::Status MergeGpuNodes(GpuNode* src, GpuNode* dst) {
    for (int j = 1; j < src->inputs.size(); ++j) {
      dst->inputs.push_back(src->inputs[j]);
    }
    dst->outputs[0] = src->outputs[0];
    dst->name += " linked : " + src->name;
    return dst->gpu_operation->AddOperation(src->gpu_operation.get());
  }
  
  flatbuffers::Offset<data::TensorDescWithId> Encode(
      const TensorDescriptor& desc, const ValueId& id,
      flatbuffers::FlatBufferBuilder* builder) {
    auto desc_fb = Encode(desc, builder);
    data::TensorDescWithIdBuilder desc_builder(*builder);
    desc_builder.add_desc(desc_fb);
    desc_builder.add_id(id);
    return desc_builder.Finish();
  }
  
  flatbuffers::Offset<data::GpuNode> Encode(
      const GpuNode& node, flatbuffers::FlatBufferBuilder* builder) {
    auto op_fb = Encode(*node.gpu_operation, builder);
    std::vector<int32_t> in_ids(node.inputs.size());
    for (int i = 0; i < in_ids.size(); ++i) {
      in_ids[i] = node.inputs[i];
    }
    std::vector<int32_t> out_ids(node.outputs.size());
    for (int i = 0; i < out_ids.size(); ++i) {
      out_ids[i] = node.outputs[i];
    }
    auto in_ids_fb = builder->CreateVector(in_ids);
    auto out_ids_fb = builder->CreateVector(out_ids);
    auto name_fb = builder->CreateString(node.name);
    data::GpuNodeBuilder node_builder(*builder);
    node_builder.add_gpu_op(op_fb);
    node_builder.add_input_ids(in_ids_fb);
    node_builder.add_output_ids(out_ids_fb);
    node_builder.add_name(name_fb);
    return node_builder.Finish();
  }
  
  absl::Status Decode(const data::GpuNode* fb_node, GpuNode* node) {
    GPUOperation op;
    RETURN_IF_ERROR(Decode(fb_node->gpu_op(), &op));
    node->gpu_operation = absl::make_unique<GPUOperation>(std::move(op));
    for (auto in_fb : *fb_node->input_ids()) {
      node->inputs.push_back(in_fb);
    }
    for (auto out_fb : *fb_node->output_ids()) {
      node->outputs.push_back(out_fb);
    }
    node->name = std::string(fb_node->name()->c_str(), fb_node->name()->size());
  
    return absl::OkStatus();
  }
  
  bool IsAssociativeLinkableOp(const Node& node,
                               const std::vector<Value*>& inputs,
                               const std::vector<Value*>& outputs) {
    if (inputs.size() == 1) {
      return false;
    }
    const OperationType op_type = OperationTypeFromString(node.operation.type);
    if (op_type != OperationType::ADD && op_type != OperationType::MUL) {
      return false;
    }
  
    const auto dst_shape = outputs[0]->tensor.shape;
    for (int i = 0; i < inputs.size(); ++i) {
      const auto src_shape = inputs[i]->tensor.shape;
      if (dst_shape.b != src_shape.b && src_shape.b == 1) {
        return false;
      }
      if (dst_shape.h != src_shape.h && src_shape.h == 1) {
        return false;
      }
      if (dst_shape.w != src_shape.w && src_shape.w == 1) {
        return false;
      }
      if (dst_shape.c != src_shape.c && src_shape.c == 1) {
        return false;
      }
    }
    return true;
  }
  
  absl::Status CheckExternalTensorDescription(const GpuInfo& gpu_info,
                                              const TensorDescriptor& tensor_desc,
                                              const BHWC& shape,
                                              DataType data_type) {
    if (tensor_desc.data_type != data_type) {
      return absl::InvalidArgumentError(
          "Global precision and precision of predefined/external tensors must be "
          "synchronized.");
    }
    const bool tensor_supported_layout = tensor_desc.layout == Layout::HWDC ||
                                         tensor_desc.layout == Layout::BHWDC ||
                                         tensor_desc.layout == Layout::HWC ||
                                         tensor_desc.layout == Layout::BHWC;
    if (!tensor_supported_layout) {
      return absl::InvalidArgumentError(
          "Currently no support of this layouts for spatial tensors.");
    }
    const bool has_depth =
        tensor_desc.layout == Layout::HWDC || tensor_desc.layout == Layout::BHWDC;
    if (has_depth) {
      return absl::InvalidArgumentError(
          "Currently no support of Depth dimension in predefined/external "
          "tensors.");
    }
    const bool has_batch =
        tensor_desc.layout == Layout::BHWC || tensor_desc.layout == Layout::BHWDC;
    if (has_batch && shape.b == 1) {
      return absl::InvalidArgumentError("Wrong layout, batch mismatch.");
    }
    if (!has_batch && shape.b != 1) {
      return absl::InvalidArgumentError("Wrong layout, batch mismatch.");
    }
    if (!CanCreateTensorWithShape(gpu_info, shape, tensor_desc).ok()) {
      return absl::UnavailableError(
          "Current device can not allocate tensor with this shape for "
          "predefined/external descriptor.");
    }
    return absl::OkStatus();
  }
  
  // Helper class for creating descriptors for appropriate tensors from
  // GraphFloat32
  // Also allows to create descriptors for new tensors(not present in
  // GraphFloat32)
  class TensorReserver {
   public:
    TensorReserver() : next_(0) {}
    ValueId Add(const TensorDescriptor& dummy) {
      reservations_[next_] = dummy;
      return next_++;
    }
    void Add(ValueId id, const TensorDescriptor& dummy) {
      reservations_[id] = dummy;
    }
    void SetNext(ValueId id) { next_ = id; }
    TensorDescriptor Get(ValueId id) { return reservations_[id]; }
  
   public:
    absl::flat_hash_map<ValueId, TensorDescriptor> reservations_;
    ValueId next_;
  };
  
  absl::Status ReserveGraphTensors(const CreateGpuModelInfo& create_info,
                                   const GpuInfo& gpu_info,
                                   const GraphFloat32& graph,
                                   TensorReserver* tensor_reserver) {
    ValueId max_id = 0;
    auto tensors = graph.values();
  
    //author:fu
    TFLITE_LOG(INFO) << "graph tensor size = " << tensors.size() << std::endl << "values idx:";
    for (auto& t : tensors){
      TFLITE_LOG(INFO) << t->id << " ";
    }
    auto data_type = DeduceDataTypeFromPrecision(create_info.precision);
    for (auto& t : tensors) {
      const auto shape = graph.GetValue(t->id)->tensor.shape;
      auto it_predefined = create_info.predefined.find(t->id);
      auto it_immutable_external =
          create_info.external_immutable_tensors.find(t->id);
      auto it_mutable_external = create_info.external_mutable_tensors.find(t->id);
      int external_categories_count = 0;
      TensorDescriptor tensor_desc;
      if (it_predefined != create_info.predefined.end()) {
        external_categories_count++;
        tensor_desc = it_predefined->second;
      }
      if (it_immutable_external != create_info.external_immutable_tensors.end()) {
        external_categories_count++;
        tensor_desc = it_immutable_external->second->GetDescriptor();
      }
      if (it_mutable_external != create_info.external_mutable_tensors.end()) {
        external_categories_count++;
        tensor_desc = it_mutable_external->second;
      }
      if (external_categories_count > 1) {
        return absl::InvalidArgumentError(
            "Tensors ids from predefined / external_immutable_tensors / "
            "external_mutable_tensors should not intersect.");
      }
      if (external_categories_count == 1) {
        if (!(graph.IsGraphInput(t->id) || graph.IsGraphOutput(t->id))) {
          return absl::InvalidArgumentError(
              "Currently external can be used only for graph inputs/outputs");
        }
        RETURN_IF_ERROR(CheckExternalTensorDescription(gpu_info, tensor_desc,
                                                       shape, data_type));
      } else {
        TensorStorageType storage_type = create_info.storage_type;
        Layout layout = shape.b == 1 ? Layout::HWC : Layout::BHWC;
        if (graph.IsGraphInput(t->id) || graph.IsGraphOutput(t->id)) {
          if (shape.c < 4 &&
              CanCreateTensorWithShape(
                  gpu_info, shape,
                  TensorDescriptor{data_type,
                                   TensorStorageType::SINGLE_TEXTURE_2D, layout})
                  .ok()) {
            storage_type = TensorStorageType::SINGLE_TEXTURE_2D;
          }
        }
        RETURN_IF_ERROR(SelectBestStorageType(gpu_info, shape, storage_type,
                                              data_type, layout, &storage_type));
        tensor_desc = TensorDescriptor{data_type, storage_type, layout};
        if (gpu_info.IsApiMetal() &&
            storage_type == TensorStorageType::TEXTURE_2D) {
          tensor_desc.use_buffer_for_write_only_2d_texture = true;
        }
      }
      tensor_desc.shape = BHWDC(shape.b, shape.h, shape.w, 1, shape.c);
      tensor_reserver->Add(t->id, tensor_desc);
      max_id = std::max(max_id, t->id);
    }
    tensor_reserver->SetNext(max_id + 1);
    return absl::OkStatus();
  }
  
  absl::Status ConvertOperations(const GpuInfo& gpu_info,
                                 const GraphFloat32& graph,
                                 const CreateGpuModelInfo& create_info,
                                 TensorReserver* tensor_reserver,
                                 GpuModel* gpu_model) {
    std::map<ValueId, TensorDescriptor> tensor_descriptors;
    const auto values = graph.values();
    for (auto value : values) {
      tensor_descriptors[value->id] = tensor_reserver->Get(value->id);
    }
    std::set<NodeId> consumed_nodes;
    std::vector<Node*> graph_nodes = graph.nodes();
    std::map<ValueId, int>
        tensor_usages;  // keeps latest index of operation that updated tensor
    for (const auto& input : gpu_model->input_ids_and_refs) {
      tensor_usages[input.first] = -1;  // so as inputs "updated" before operation
                                        // 0, we will mark them with -1
    }
    for (int i = 0; i < graph_nodes.size(); ++i) {
      const Node& node = *graph_nodes[i];
      if (consumed_nodes.find(node.id) != consumed_nodes.end()) {
        continue;
      }
      auto op_type = OperationTypeFromString(node.operation.type);
      if (op_type == OperationType::CONSTANT) {
        auto attr =
            absl::any_cast<ConstTensorAttributes>(node.operation.attributes);
        auto outputs = graph.FindOutputs(node.id);
        gpu_model->const_tensors[outputs[0]->id] =
            tensor_reserver->Get(outputs[0]->id);
        gpu_model->const_tensors[outputs[0]->id].UploadData(attr.tensor);
        continue;
      }
      GPUOperationsSubgraph gpu_subgraph;
      if (create_info.hints.Check(ModelHints::kAllowSpecialKernels) &&
          GPUSubgraphFromGraph(gpu_info, create_info.precision, graph, node.id,
                               tensor_descriptors, &consumed_nodes, &gpu_subgraph)
              .ok()) {
        // Mapping of subgraph (set of nodes) to GPU operations. Should happen
        // before straigtforward mapping.
      } else {
        // Straigtforward mapping of one graph node to GPU operations.
        auto inputs = graph.FindInputs(node.id);
        auto outputs = graph.FindOutputs(node.id);
        // Reordering of input ids and updating of temporary tensors_usage struct.
        // To have better linking we need linking tensor(latest written during
        // linear execution) on first position.
        if (IsAssociativeLinkableOp(node, inputs, outputs)) {
          int latest_written_tensor_index = 0;
          int last_usage = tensor_usages[inputs[0]->id];
          for (int j = 1; j < inputs.size(); ++j) {
            if (tensor_usages[inputs[j]->id] > last_usage) {
              last_usage = tensor_usages[inputs[j]->id];
              latest_written_tensor_index = j;
            }
          }
          std::swap(inputs[0], inputs[latest_written_tensor_index]);
        }
        consumed_nodes.insert(node.id);
        OperationDef op_def;
        op_def.precision = create_info.precision;
        for (int j = 0; j < inputs.size(); ++j) {
          op_def.src_tensors.push_back(tensor_reserver->Get(inputs[j]->id));
        }
        for (int j = 0; j < outputs.size(); ++j) {
          op_def.dst_tensors.push_back(tensor_reserver->Get(outputs[j]->id));
        }
        RETURN_IF_ERROR(GPUOperationFromNode(gpu_info, op_def, create_info.hints,
                                             inputs, outputs, node,
                                             &gpu_subgraph));
      }
      absl::flat_hash_map<int, ValueId> mapping_to_global_ids;
      for (int j = 0; j < gpu_subgraph.new_tensors.size(); ++j) {
        const auto& t = gpu_subgraph.new_tensors[j];
        TensorDescriptor td = t.second;
        td.shape = BHWDC(t.first.b, t.first.h, t.first.w, 1, t.first.c);
        auto global_id = tensor_reserver->Add(td);
        mapping_to_global_ids[j] = global_id;
      }
      for (auto& gpu_op : gpu_subgraph.operations) {
        GpuNode gpu_node;
        gpu_node.gpu_operation = std::move(gpu_op.operation);
        gpu_node.inputs.resize(gpu_op.input_ids.size());
        for (int j = 0; j < gpu_op.input_ids.size(); ++j) {
          int id = gpu_op.input_ids[j];
          if (id >= 0) {
            gpu_node.inputs[j] = id;
          } else {
            gpu_node.inputs[j] = mapping_to_global_ids[-(id + 1)];
          }
        }
        gpu_node.outputs.resize(gpu_op.output_ids.size());
        for (int j = 0; j < gpu_op.output_ids.size(); ++j) {
          int id = gpu_op.output_ids[j];
          if (id >= 0) {
            gpu_node.outputs[j] = id;
            tensor_usages[id] = i;
          } else {
            gpu_node.outputs[j] = mapping_to_global_ids[-(id + 1)];
          }
        }
        gpu_node.name = gpu_op.name;
        gpu_model->nodes.push_back(std::move(gpu_node));
      }
    }
  
    return absl::OkStatus();
  }
  
  absl::Status MergeNodes(GpuModel* gpu_model) {
    absl::flat_hash_set<ValueId> ready_tensors;
    for (const auto& input : gpu_model->input_ids_and_refs) {
      ready_tensors.insert(input.first);
    }
    auto& nodes = gpu_model->nodes;
    for (int i = 0; i < nodes.size(); ++i) {
      auto& node = nodes[i];
      for (const auto& out_id : node.outputs) {
        ready_tensors.insert(out_id);
      }
      if (node.outputs.size() != 1) {
        continue;
      }
      std::vector<int> next_nodes;
      int link_index = 0;
      for (int j = i + 1; j < nodes.size(); ++j) {
        for (int k = 0; k < nodes[j].inputs.size(); ++k) {
          if (nodes[j].inputs[k] == node.outputs[0]) {
            next_nodes.push_back(j);
            link_index = k;
          }
        }
      }
      if (next_nodes.size() != 1 || link_index != 0) {
        continue;
      }
      auto& linkable_node = nodes[next_nodes[0]];
      if (!linkable_node.gpu_operation->IsLinkable() ||
          linkable_node.outputs.size() != 1 ||
          !IsReady(ready_tensors, linkable_node)) {
        continue;
      }
      const auto& original_dst_def =
          node.gpu_operation->GetDefinition().dst_tensors[0];
      const auto& link_dst_def =
          linkable_node.gpu_operation->GetDefinition().dst_tensors[0];
      if (original_dst_def != link_dst_def) {
        continue;
      }
      RETURN_IF_ERROR(MergeGpuNodes(&linkable_node, &node));
      nodes.erase(nodes.begin() + next_nodes[0]);
      i -= 1;
    }
    return absl::OkStatus();
  }
  
  void CopyExternals(const GraphFloat32& graph, GpuModel* gpu_model) {
    const auto inputs = graph.inputs();
    TFLITE_LOG(INFO) << "gpu_model inputs: " << std::endl;
  
    for (const auto& value : inputs) {
      gpu_model->input_ids_and_refs.push_back({value->id, value->tensor.ref});
      TFLITE_LOG(INFO) << "value_id=" << value->id << "value->tensor.ref" << value->tensor.ref << std::endl;
    }
  
    const auto variable_inputs = graph.variable_inputs();
    for (const auto& value : variable_inputs) {
      gpu_model->variable_ids_and_refs.push_back({value->id, value->tensor.ref});
    }
    TFLITE_LOG(INFO) << "gpu_model inputs: " << std::endl;
    const auto outputs = graph.outputs();
    for (const auto& value : outputs) {
      gpu_model->output_ids_and_refs.push_back({value->id, value->tensor.ref});
      TFLITE_LOG(INFO) << "value_id=" << value->id << "value->tensor.ref" << value->tensor.ref << std::endl;
    }
  }
  
  // Removing tensors that was fused in complex operations
  void RemoveUnusedTensors(GpuModel* gpu_model) {
    absl::flat_hash_set<ValueId> used_tensors;
    for (const auto& node : gpu_model->nodes) {
      for (const auto& id : node.inputs) {
        used_tensors.insert(id);
      }
      for (const auto& id : node.outputs) {
        used_tensors.insert(id);
      }
    }
    for (auto it = gpu_model->tensors.begin(); it != gpu_model->tensors.end();) {
      if (used_tensors.find(it->first) == used_tensors.end()) {
        gpu_model->tensors.erase(it++);
      } else {
        ++it;
      }
    }
  }
  
  // Serialized model will lose polymorphic properties for GpuOperations.
  // Here we will retrieve some information needed for generic execution of
  // GpuOperations. Specifically, BindArguments and RecalculateGridSize must be
  // executed.
  absl::Status ResolvePolymorphicArgs(GpuModel* gpu_model) {
    class DummySpatialTensor : public GpuSpatialTensor {
     public:
      DummySpatialTensor() = default;
      explicit DummySpatialTensor(const BHWDC& shape,
                                  const TensorDescriptor& tensor_desc)
          : shape_(shape), tensor_desc_(tensor_desc) {}
      ~DummySpatialTensor() override = default;
  
      int Width() const override { return shape_.w; }
      int Height() const override { return shape_.h; }
      int Depth() const override { return shape_.d; }
      int Channels() const override { return shape_.c; }
      int Slices() const override { return DivideRoundUp(shape_.c, 4); }
      int Batch() const override { return shape_.b; }
  
      TensorDescriptor GetDescriptor() const override { return tensor_desc_; }
  
     private:
      BHWDC shape_;
      TensorDescriptor tensor_desc_;
    };
  
    for (auto& node : gpu_model->nodes) {
      std::vector<DummySpatialTensor> src_tensors(node.inputs.size());
      for (int i = 0; i < node.inputs.size(); ++i) {
        const auto& tensor_desc = gpu_model->tensors[node.inputs[i]];
        src_tensors[i] = DummySpatialTensor(tensor_desc.shape, tensor_desc);
        node.gpu_operation->SetSrc(&src_tensors[i], i);
      }
      std::vector<DummySpatialTensor> dst_tensors(node.outputs.size());
      for (int i = 0; i < node.outputs.size(); ++i) {
        const auto& tensor_desc = gpu_model->tensors[node.outputs[i]];
        dst_tensors[i] = DummySpatialTensor(tensor_desc.shape, tensor_desc);
        node.gpu_operation->SetDst(&dst_tensors[i], i);
      }
      RETURN_IF_ERROR(
          node.gpu_operation->BindArguments(&node.gpu_operation->args_));
      node.gpu_operation->RecalculateGridSize();
    }
    return absl::OkStatus();
  }
  
  }  // namespace
  
  absl::Status GraphToGpuModel(const GraphFloat32& graph,
                               const CreateGpuModelInfo& create_info,
                               const GpuInfo& gpu_info, GpuModel* gpu_model) {
    TensorReserver tensor_reserver;
    RETURN_IF_ERROR(
        ReserveGraphTensors(create_info, gpu_info, graph, &tensor_reserver));
    CopyExternals(graph, gpu_model);
    RETURN_IF_ERROR(ConvertOperations(gpu_info, graph, create_info,
                                      &tensor_reserver, gpu_model));
    RETURN_IF_ERROR(MergeNodes(gpu_model));
    gpu_model->tensors = std::move(tensor_reserver.reservations_);
    RemoveUnusedTensors(gpu_model);
  
    for (auto& node : gpu_model->nodes) {
      RETURN_IF_ERROR(node.gpu_operation->AssembleCode(gpu_info));
    }
  
    return ResolvePolymorphicArgs(gpu_model);
  }
  
  flatbuffers::Offset<data::GpuModel> Encode(
      const GpuModel& gpu_model, flatbuffers::FlatBufferBuilder* builder) {
    std::vector<int32_t> in_ids(gpu_model.input_ids_and_refs.size());
    std::vector<int64_t> in_refs(gpu_model.input_ids_and_refs.size());
    for (int i = 0; i < in_ids.size(); ++i) {
      in_ids[i] = gpu_model.input_ids_and_refs[i].first;
      in_refs[i] = gpu_model.input_ids_and_refs[i].second;
    }
    auto in_ids_fb = builder->CreateVector(in_ids);
    auto in_refs_fb = builder->CreateVector(in_refs);
  
    std::vector<int32_t> out_ids(gpu_model.output_ids_and_refs.size());
    std::vector<int64_t> out_refs(gpu_model.output_ids_and_refs.size());
    for (int i = 0; i < out_ids.size(); ++i) {
      out_ids[i] = gpu_model.output_ids_and_refs[i].first;
      out_refs[i] = gpu_model.output_ids_and_refs[i].second;
    }
    auto out_ids_fb = builder->CreateVector(out_ids);
    auto out_refs_fb = builder->CreateVector(out_refs);
  
    std::vector<flatbuffers::Offset<data::GpuNode>> nodes_fb;
    for (int i = 0; i < gpu_model.nodes.size(); ++i) {
      auto node_fb = Encode(gpu_model.nodes[i], builder);
      nodes_fb.push_back(node_fb);
    }
    auto nodes_fb_vec = builder->CreateVector(nodes_fb);
  
    std::vector<flatbuffers::Offset<data::TensorDescWithId>> tensors_fb;
    for (const auto& tensor : gpu_model.tensors) {
      auto tensor_fb = Encode(tensor.second, tensor.first, builder);
      tensors_fb.push_back(tensor_fb);
    }
    auto tensors_fb_vec = builder->CreateVector(tensors_fb);
  
    std::vector<flatbuffers::Offset<data::TensorDescWithId>> const_tensors_fb;
    for (const auto& tensor : gpu_model.const_tensors) {
      auto tensor_fb = Encode(tensor.second, tensor.first, builder);
      const_tensors_fb.push_back(tensor_fb);
    }
    auto const_tensors_fb_vec = builder->CreateVector(const_tensors_fb);
  
    std::vector<flatbuffers::Offset<data::PairOfValueIds>>
        variable_ids_and_refs_fb;
    for (auto& pair : gpu_model.variable_ids_and_refs) {
      data::PairOfValueIdsBuilder pair_builder(*builder);
      pair_builder.add_first(pair.first);
      pair_builder.add_second(pair.second);
      variable_ids_and_refs_fb.push_back(pair_builder.Finish());
    }
    auto variable_ids_and_refs_fb_vec =
        builder->CreateVector(variable_ids_and_refs_fb);
  
    data::GpuModelBuilder gpu_model_builder(*builder);
    gpu_model_builder.add_nodes(nodes_fb_vec);
    gpu_model_builder.add_tensors(tensors_fb_vec);
    gpu_model_builder.add_const_tensors(const_tensors_fb_vec);
    gpu_model_builder.add_input_ids(in_ids_fb);
    gpu_model_builder.add_output_ids(out_ids_fb);
    gpu_model_builder.add_variable_ids_and_refs(variable_ids_and_refs_fb_vec);
    gpu_model_builder.add_input_refs(in_refs_fb);
    gpu_model_builder.add_output_refs(out_refs_fb);
    return gpu_model_builder.Finish();
  }
  
  absl::Status Decode(const data::GpuModel* fb_gpu_model, GpuModel* gpu_model) {
    gpu_model->nodes.resize(fb_gpu_model->nodes()->size());
    int counter = 0;
    for (auto node_fb : *fb_gpu_model->nodes()) {
      RETURN_IF_ERROR(Decode(node_fb, &gpu_model->nodes[counter]));
      counter++;
    }
  
    for (const auto& tensor_fb : *fb_gpu_model->tensors()) {
      TensorDescriptor desc;
      Decode(tensor_fb->desc(), &desc);
      gpu_model->tensors[tensor_fb->id()] = std::move(desc);
    }
    for (const auto& tensor_fb : *fb_gpu_model->const_tensors()) {
      TensorDescriptor desc;
      Decode(tensor_fb->desc(), &desc);
      gpu_model->const_tensors[tensor_fb->id()] = std::move(desc);
    }
    for (int i = 0; i < fb_gpu_model->input_ids()->size(); ++i) {
      gpu_model->input_ids_and_refs.push_back(
          {(*fb_gpu_model->input_ids())[i], (*fb_gpu_model->input_refs())[i]});
    }
    for (int i = 0; i < fb_gpu_model->output_ids()->size(); ++i) {
      gpu_model->output_ids_and_refs.push_back(
          {(*fb_gpu_model->output_ids())[i], (*fb_gpu_model->output_refs())[i]});
    }
  
    for (auto variable_id : *fb_gpu_model->variable_ids_and_refs()) {
      gpu_model->variable_ids_and_refs.push_back(
          {variable_id->first(), variable_id->second()});
    }
    return absl::OkStatus();
  }
  
  absl::Status RunGraphTransformsForGpuModel(GraphFloat32* graph) {
    auto merge_padding_transform = NewMergePaddingWithAdd();
    auto add_bias_transform = NewAddBias();
    auto pooling_to_reduce_op = NewGlobalPoolingToReduceOp();
    ModelTransformer transformer(graph);
    if (!transformer.Apply("add_bias", add_bias_transform.get())) {
      return absl::InternalError("Invalid add_bias transform");
    }
    if (!transformer.Apply("merge_padding", merge_padding_transform.get())) {
      return absl::InternalError("Invalid merge_padding transform");
    }
    if (!transformer.Apply("global pooling to mean",
                           pooling_to_reduce_op.get())) {
      return absl::InternalError("Invalid global pooling to mean transform");
    }
    return absl::OkStatus();
  }
  
  }  // namespace gpu
  }  // namespace tflite
  