| 1169 | // |
| 1170 | template <typename OutputType> |
| 1171 | void CopyDataToOutput(const absl::Span<complex128> data, int64 output_start, |
| 1172 | int64 fft_rank, FftType fft_type, |
| 1173 | const absl::Span<const int64> fft_lengths, |
| 1174 | const absl::Span<const int64> fft_strides, |
| 1175 | const absl::Span<const int64> output_lengths, |
| 1176 | const absl::Span<const int64> output_strides, |
| 1177 | Literal* output_literal) { |
| 1178 | const bool output_is_truncated = fft_type == FftType::RFFT; |
| 1179 | |
| 1180 | // Base case for recursive copy of the results to the output. The code avoids |
| 1181 | // making a recursive call for each output element by handling axis 0 in the |
| 1182 | // loop (as opposed to making "axis < 0" to be the base case). |
| 1183 | OutputType* output_data = output_literal->data<OutputType>().data(); |
| 1184 | auto base_case = [&](int64 axis, int64 dst_index, int64 src_index, |
| 1185 | bool within_src_bounds) { |
| 1186 | if (axis == 0) { |
| 1187 | // Drop negative frequencies for RFFT. |
| 1188 | const int64 length = fft_lengths[axis]; |
| 1189 | const int64 ub = output_is_truncated ? (length / 2) + 1 : length; |
| 1190 | for (int64 i = 0; i < output_lengths[axis]; i++) { |
| 1191 | OutputType value = OutputType(0); |
| 1192 | // Read data only if the index is within bounds. |
| 1193 | if (within_src_bounds && i < ub) { |
| 1194 | value = GetAs<OutputType, complex128>( |
| 1195 | data[src_index + i * fft_strides[axis]]); |
| 1196 | } |
| 1197 | output_data[dst_index + i * output_strides[axis]] = value; |
| 1198 | } |
| 1199 | return true; |
| 1200 | } |
| 1201 | return false; |
| 1202 | }; |
| 1203 | GenerateIndices(output_lengths, output_strides, fft_lengths, fft_strides, |
| 1204 | fft_rank, output_start, 0, base_case); |
| 1205 | } |
| 1206 | |
| 1207 | // Determine the type to use with the CopyDataFromInput<> template above. |
| 1208 | bool CopyDataFromInput(const Literal& input_literal, int64 input_start, |
no test coverage detected