Layout is Pc,H,W,C4 where P - is a plane based on channels.
| 224 | |
| 225 | // Layout is Pc,H,W,C4 where P - is a plane based on channels. |
| 226 | Status ConvertToPHWC4(absl::Span<const float> in, const BHWC& shape, |
| 227 | absl::Span<float> out) { |
| 228 | RETURN_IF_ERROR(ValidateConvertToPHWC4(in, shape, out)); |
| 229 | if (shape.c == 4) { |
| 230 | std::memcpy(out.data(), in.data(), |
| 231 | shape.DimensionsProduct() * sizeof(float)); |
| 232 | return OkStatus(); |
| 233 | } |
| 234 | // Layout is Pc,H,W,C4 where P - is a plane based on channels. |
| 235 | int num_planes = IntegralDivideRoundUp(shape.c, kPhwc4ChannelsInPlane); |
| 236 | const int num_pixels = shape.h * shape.w; |
| 237 | // A layer is a set of kPhwc4ChannelsInPlane channels images. |
| 238 | const int num_full_planes = shape.c / kPhwc4ChannelsInPlane; |
| 239 | for (int b = 0; b < shape.b; b++) { |
| 240 | float* dest = |
| 241 | out.data() + b * num_pixels * num_planes * kPhwc4ChannelsInPlane; |
| 242 | for (int p = 0; p < num_full_planes; p++) { |
| 243 | const float* src = |
| 244 | in.data() + shape.LinearIndex({b, 0, 0, p * kPhwc4ChannelsInPlane}); |
| 245 | for (int i = 0; i < num_pixels; i++) { |
| 246 | std::memcpy(dest, src, kPhwc4ChannelsInPlane * sizeof(float)); |
| 247 | src += shape.c; |
| 248 | dest += kPhwc4ChannelsInPlane; |
| 249 | } |
| 250 | } |
| 251 | } |
| 252 | |
| 253 | // Padding last kPhwc4ChannelsInPlane-channel layer to multiple of |
| 254 | // kPhwc4ChannelsInPlane. |
| 255 | const int padded_size = num_pixels * num_planes * kPhwc4ChannelsInPlane; |
| 256 | const int remaining_channels = |
| 257 | shape.c - num_full_planes * kPhwc4ChannelsInPlane; |
| 258 | if (remaining_channels == 0) { |
| 259 | return OkStatus(); |
| 260 | } |
| 261 | for (int b = 0; b < shape.b; b++) { |
| 262 | const float* src = |
| 263 | in.data() + |
| 264 | shape.LinearIndex({b, 0, 0, num_full_planes * kPhwc4ChannelsInPlane}); |
| 265 | float* dest = out.data() + b * padded_size + |
| 266 | num_pixels * num_full_planes * kPhwc4ChannelsInPlane; |
| 267 | for (int p = 0; p < num_pixels; p++) { |
| 268 | std::memcpy(dest, src, remaining_channels * sizeof(float)); |
| 269 | std::memset(dest + remaining_channels, 0, |
| 270 | (4 - remaining_channels) * sizeof(float)); |
| 271 | src += shape.c; |
| 272 | dest += kPhwc4ChannelsInPlane; |
| 273 | } |
| 274 | } |
| 275 | return OkStatus(); |
| 276 | } |
| 277 | |
| 278 | // Layout is Pc,H,W,C4 where P - is a plane based on channels. |
| 279 | Status ConvertToPHWC4Half(absl::Span<const float> in, const BHWC& shape, |