Function GetKernelForConv

tensorflow/lite/delegates/gpu/metal/kernels/conv.cc:151–237

}

std::string GetKernelForConv(const Convolution2DAttributes& params) {
  const int num_output_slices = GetNumOutputSlices(params.weights.shape.o);
  std::string code;
  code.reserve(16 * 1024);  // Reserve large enough buffer.
  const bool is_1x1 =
      params.weights.shape.w == 1 && params.weights.shape.h == 1;
  const bool is_strided = params.strides.w > 1 || params.strides.h > 1;
  const int src_group_size = GetSrcBatchSize(params.weights.shape.o);

  const int src_depth = IntegralDivideRoundUp(params.weights.shape.i, 4);
  const int src_groups = src_depth / src_group_size;
  const int src_depth_aligned = AlignByN(src_depth, src_group_size);
  const int reminder_src_depth = src_depth - src_groups * src_group_size;

  code = absl::Substitute(R"(
    #include <metal_stdlib>
    using namespace metal;
    constant int src_depth_groups = $0;
    constant int src_offset = $1;
    constant int kernel_x = $2;
    constant int kernel_y = $3;
    struct uniforms {
      int4 stride_padding;
      int4 dillation_layer_offsets;
      int4 size;
      int4 z_offset;
    };
    $$0
    kernel void ComputeFunction(
                                $$1
                                uint tid[[thread_index_in_threadgroup]],
                                uint3 gid[[thread_position_in_grid]])
      {
        const bool outside = static_cast<int>(gid.x) >= params.size.z ||
          static_cast<int>(gid.y) >= params.size.w;
  )",
                          src_groups, src_depth_aligned, params.weights.shape.w,
                          params.weights.shape.h);
  code += GetValuesDeclarationPart(num_output_slices, is_1x1);

  if (!is_1x1) {
    code += R"(
      for(int ky = 0; ky < kernel_y; ++ky) {
        for(int kx = 0; kx < kernel_x; ++kx) {
          int2 coords = int2(gid.xy) * params.stride_padding.xy + int2(kx, ky) *
            params.dillation_layer_offsets.xy - params.stride_padding.zw;
          const bool el_outside = coords.x < 0 || coords.y < 0 || coords.x >= params.size.x ||
            coords.y >= params.size.y;
          const FLT multiplier = el_outside ? 0.0f : 1.0f;
    )";
  } else {
    code += "const FLT multiplier = 1.0f;\n";
    code += "int2 coords = int2(gid.xy)";
    if (is_strided) {
      code += " * params.stride_padding.xy";
    }
    code += ";\n";
  }