Function Convolution1x1

tensorflow/lite/delegates/gpu/metal/kernels/conv.cc:970–1026

}

std::vector<ComputeTaskDescriptorPtr> Convolution1x1(
    int id, ValueId input_id, ValueId output_id,
    const Convolution2DAttributes& params,
    const metal::RuntimeOptions& options) {
  auto desc = std::make_shared<ComputeTaskDescriptor>();
  desc->id = id;
  desc->is_linkable = false;
  const int z_out = GetNumOutputSlices(params.weights.shape.o);
  desc->shader_source = GetKernelForConv1x1(params, z_out);

  desc->input_buffers = {
      {input_id, "device FLT4* const src_buffer"},
  };

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

  auto weights_reordered = ReorderWeightsForConv(params, z_out);
  auto weights =
      options.storage_precision == metal::RuntimeOptions::Precision::FP32
          ? VectorToUint8Vector(weights_reordered)
          : VectorFloatToHalf(weights_reordered);
  auto biases =
      options.storage_precision == metal::RuntimeOptions::Precision::FP32
          ? VectorToUint8Vector(params.bias.data)
          : VectorFloatToHalf(params.bias.data);
  desc->immutable_buffers = {
      {"device FLT4* const filters", weights},
      {"device FLT4* const biases", biases},
  };

  desc->uniform_buffers = {
      {"constant uniforms& params",
       [input_id, output_id, params](const std::map<ValueId, BHWC>& buffers) {
         const auto& input_dimensions = buffers.find(input_id)->second;
         const auto& output_dimensions = buffers.find(output_id)->second;
         return GetUniformBufferForConv(input_dimensions, output_dimensions,
                                        params);
       }},
  };

  desc->resize_function = [output_id,
                           params](const std::map<ValueId, BHWC>& buffers) {
    const auto& output_dims = buffers.find(output_id)->second;
    const uint3 group_size = GetWorkGroupForConv();
    const uint3 groups_count = GetGroupsCountForConv(group_size, output_dims);
    return std::make_pair(
        group_size, uint3{groups_count.z, groups_count.x, groups_count.y});
  };

  return {desc};
}