| 165 | } |
| 166 | |
| 167 | Status ConvertToPIOHW4(absl::Span<const float> in, const OHWI& shape, |
| 168 | absl::Span<float> out) { |
| 169 | if (in.size() != shape.DimensionsProduct()) { |
| 170 | return InvalidArgumentError(absl::StrCat( |
| 171 | "ConvertToPIOHW4: Input data size does not match expected size: ", |
| 172 | in.size(), " != ", shape.DimensionsProduct())); |
| 173 | } |
| 174 | if (out.size() != GetElementsSizeForPIOHW4(shape)) { |
| 175 | return InvalidArgumentError(absl::StrCat( |
| 176 | "ConvertToPIOHW4: Output data size does not match expected size: ", |
| 177 | out.size(), " != ", GetElementsSizeForPIOHW4(shape))); |
| 178 | } |
| 179 | |
| 180 | int32_t output_channels = shape.o * shape.i; |
| 181 | int32_t num_planes = |
| 182 | IntegralDivideRoundUp(output_channels, kPiohw4ChannelsInPlane); |
| 183 | float* output = out.data(); |
| 184 | for (int p = 0; p < num_planes; ++p) { |
| 185 | for (int h = 0; h < shape.h; ++h) { |
| 186 | for (int w = 0; w < shape.w; ++w) { |
| 187 | for (int c = 0; c < kPiohw4ChannelsInPlane; ++c) { |
| 188 | int output_c = p * kPiohw4ChannelsInPlane + c; |
| 189 | (*output++) = output_c >= output_channels |
| 190 | ? 0 |
| 191 | : in[shape.LinearIndex({output_c % shape.o, h, w, |
| 192 | output_c / shape.o})]; |
| 193 | } |
| 194 | } |
| 195 | } |
| 196 | } |
| 197 | return OkStatus(); |
| 198 | } |
| 199 | |
| 200 | std::vector<float> ConvertToPIOHW4( |
| 201 | const Tensor<OHWI, DataType::FLOAT32>& tensor) { |
no test coverage detected