Function ConvertToPHWC4Half

tensorflow/lite/delegates/gpu/common/convert.cc:279–357

Layout is Pc,H,W,C4 where P - is a plane based on channels.

// Layout is Pc,H,W,C4 where P - is a plane based on channels.
Status ConvertToPHWC4Half(absl::Span<const float> in, const BHWC& shape,
                          absl::Span<HalfBits> out) {
  RETURN_IF_ERROR(ValidateConvertToPHWC4(in, shape, out));

  // Layout is Pc,H,W,C4 where P - is a plane based on channels.
  int num_planes = IntegralDivideRoundUp(shape.c, kPhwc4ChannelsInPlane);
  const int num_pixels = shape.h * shape.w;
  // A layer is a set of kPhwc4ChannelsInPlane channels images.
  const int num_full_planes = shape.c / kPhwc4ChannelsInPlane;
  for (int b = 0; b < shape.b; b++) {
    HalfBits* dest =
        out.data() + b * num_pixels * num_planes * kPhwc4ChannelsInPlane;
    for (int p = 0; p < num_full_planes; p++) {
      const float* src =
          in.data() + shape.LinearIndex({b, 0, 0, p * kPhwc4ChannelsInPlane});
      for (int i = 0; i < num_pixels; i++) {
        dest[0] = fp16_ieee_from_fp32_value(src[0]);
        dest[1] = fp16_ieee_from_fp32_value(src[1]);
        dest[2] = fp16_ieee_from_fp32_value(src[2]);
        dest[3] = fp16_ieee_from_fp32_value(src[3]);
        src += shape.c;
        dest += kPhwc4ChannelsInPlane;
      }
    }
  }

  // Padding last kPhwc4ChannelsInPlane-channel layer to multiple of
  // kPhwc4ChannelsInPlane.
  const int padded_size = num_pixels * num_planes * kPhwc4ChannelsInPlane;
  const int remaining_channels =
      shape.c - num_full_planes * kPhwc4ChannelsInPlane;
  if (remaining_channels == 0) {
    return OkStatus();
  }

  for (int b = 0; b < shape.b; b++) {
    const float* src =
        in.data() +
        shape.LinearIndex({b, 0, 0, num_full_planes * kPhwc4ChannelsInPlane});
    HalfBits* dest = out.data() + b * padded_size +
                     num_pixels * num_full_planes * kPhwc4ChannelsInPlane;
    switch (remaining_channels) {
      case 1:
        for (int p = 0; p < num_pixels; p++) {
          dest[0] = fp16_ieee_from_fp32_value(src[0]);
          dest[1] = 0;
          dest[2] = 0;
          dest[3] = 0;
          src += shape.c;
          dest += kPhwc4ChannelsInPlane;
        }
        break;
      case 2:
        for (int p = 0; p < num_pixels; p++) {
          dest[0] = fp16_ieee_from_fp32_value(src[0]);
          dest[1] = fp16_ieee_from_fp32_value(src[1]);
          dest[2] = 0;
          dest[3] = 0;