Run a single step of an inner loop while running reduction, which applies the user-provided computation on the accumulator and the output element (until the reduction is completed, the output element is also used as an accumulator).
| 2276 | // (until the reduction is completed, the output element is also used as |
| 2277 | // an accumulator). |
| 2278 | static StatusOr<bool> PerformReductionStep( |
| 2279 | absl::Span<const int64> input_index, absl::Span<const int64> output_index, |
| 2280 | absl::Span<const Literal* const> input_args, absl::Span<Literal> results, |
| 2281 | HloComputation* computation, HloEvaluator* embedded_evaluator) { |
| 2282 | int num_args = results.size(); |
| 2283 | bool is_tuple = num_args > 1; |
| 2284 | |
| 2285 | absl::InlinedVector<Literal, 1> arg_values; |
| 2286 | arg_values.reserve(num_args); |
| 2287 | absl::InlinedVector<Literal, 1> accumulators; |
| 2288 | accumulators.reserve(num_args); |
| 2289 | for (int64 i = 0; i < num_args; ++i) { |
| 2290 | arg_values.emplace_back( |
| 2291 | ShapeUtil::MakeShape(input_args[i]->shape().element_type(), {})); |
| 2292 | accumulators.emplace_back( |
| 2293 | ShapeUtil::MakeShape(input_args[i]->shape().element_type(), {})); |
| 2294 | |
| 2295 | TF_RETURN_IF_ERROR( |
| 2296 | arg_values[i].CopyElementFrom(*input_args[i], input_index, {})); |
| 2297 | TF_RETURN_IF_ERROR( |
| 2298 | accumulators[i].CopyElementFrom(results[i], output_index, {})); |
| 2299 | } |
| 2300 | |
| 2301 | // Evaluate computation with specified literal operands. |
| 2302 | absl::InlinedVector<Literal*, 2> embedded_operands; |
| 2303 | for (Literal& accumulator : accumulators) { |
| 2304 | embedded_operands.push_back(&accumulator); |
| 2305 | } |
| 2306 | for (Literal& local_input : arg_values) { |
| 2307 | embedded_operands.push_back(&local_input); |
| 2308 | } |
| 2309 | |
| 2310 | TF_ASSIGN_OR_RETURN( |
| 2311 | Literal computed_result, |
| 2312 | embedded_evaluator->Evaluate(*computation, embedded_operands)); |
| 2313 | |
| 2314 | // Clear visit states so that we can use the evaluator again on the same |
| 2315 | // computation. |
| 2316 | embedded_evaluator->ResetVisitStates(); |
| 2317 | |
| 2318 | if (is_tuple) { |
| 2319 | std::vector<Literal> computed_results = computed_result.DecomposeTuple(); |
| 2320 | for (int64 i = 0; i < num_args; ++i) { |
| 2321 | TF_RETURN_IF_ERROR( |
| 2322 | results[i].CopyElementFrom(computed_results[i], {}, output_index)); |
| 2323 | } |
| 2324 | } else { |
| 2325 | TF_RETURN_IF_ERROR( |
| 2326 | results[0].CopyElementFrom(computed_result, {}, output_index)); |
| 2327 | } |
| 2328 | |
| 2329 | return true; |
| 2330 | } |
| 2331 | |
| 2332 | static StatusOr<bool> GenerateReduceOutputElement( |
| 2333 | absl::Span<const int64> output_index, |
no test coverage detected