Function Add

tensorflow/lite/delegates/gpu/metal/kernels/add.cc:69–123

}  // namespace

std::vector<ComputeTaskDescriptorPtr> Add(int id,
                                          const std::vector<ValueId> input_ids,
                                          ValueId output_id,
                                          const AddAttributes& attr,
                                          const RuntimeOptions& options) {
  auto desc = std::make_shared<ComputeTaskDescriptor>();
  desc->id = id;

  // Add scalar
  const float* add_value = absl::get_if<float>(&attr.param);
  if (add_value) {
    desc->is_linkable = true;
    desc->shader_source =
        R"(FLT4 linkable$0(FLT4 value, int linear_index, uint3 gid) {
      return value + )" +
        std::to_string(*add_value) + ";}";
    desc->input_buffers = {{input_ids[0]}};
    desc->output_buffer = {output_id};
    return {desc};
  }

  desc->is_linkable = false;
  desc->shader_source = GetAddTableCode(input_ids.size());

  for (int i = 0; i < input_ids.size(); ++i) {
    const std::string buffer_name =
        "device FLT4* const src_buffer" + std::to_string(i);
    desc->input_buffers.push_back({input_ids[i], buffer_name});
  }

  desc->output_buffer = {output_id, "device FLT4* dst_buffer",
                         [input_ids](const std::map<ValueId, BHWC>& buffers) {
                           return buffers.find(input_ids[0])->second;
                         }};

  desc->uniform_buffers = {
      {"constant uniforms& params",
       [input_ids](const std::map<ValueId, BHWC>& buffers) {
         const auto& dimension = buffers.find(input_ids[0])->second;
         std::vector<int> uniform_params = {dimension.w, dimension.h, 0, 0};
         return VectorToUint8Vector(uniform_params);
       }},
  };

  desc->resize_function = [input_ids](const std::map<ValueId, BHWC>& buffers) {
    const auto& src_dim = buffers.find(input_ids[0])->second;
    const uint3 groups_size{16, 16, 1};
    int groups_x = IntegralDivideRoundUp(src_dim.w, groups_size.x);
    int groups_y = IntegralDivideRoundUp(src_dim.h, groups_size.y);
    const int dst_layers = IntegralDivideRoundUp(src_dim.c, 4);
    int groups_z = IntegralDivideRoundUp(dst_layers, groups_size.z);
    return std::make_pair(groups_size, uint3{groups_x, groups_y, groups_z});
  };
  return {desc};
}

}  // namespace metal
}  // namespace gpu