| 3554 | } |
| 3555 | |
| 3556 | Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* hlo) { |
| 3557 | HloReduceInstruction* reduce = Cast<HloReduceInstruction>(hlo); |
| 3558 | bool multi_output_reduce = reduce->shape().IsTuple(); |
| 3559 | // For tuple reduce, we require all reduce shapes to be the same, up to the |
| 3560 | // element types, so we can just the first operand and the first result as a |
| 3561 | // representative. |
| 3562 | auto arg = reduce->inputs()[0]; |
| 3563 | auto init_value = reduce->init_values()[0]; |
| 3564 | Shape& reduce_result_shape = const_cast<Shape&>( |
| 3565 | multi_output_reduce ? reduce->shape().tuple_shapes(0) : reduce->shape()); |
| 3566 | |
| 3567 | absl::Span<const int64> dimensions(reduce->dimensions()); |
| 3568 | HloComputation* function = reduce->to_apply(); |
| 3569 | if (ShapeUtil::IsZeroElementArray(arg->shape()) || |
| 3570 | ShapeUtil::IsZeroElementArray(reduce_result_shape)) { |
| 3571 | if (multi_output_reduce) { |
| 3572 | std::vector<HloInstruction*> broadcast_inits; |
| 3573 | int64 inputs = reduce->input_count(); |
| 3574 | for (int64 i = 0; i < inputs; ++i) { |
| 3575 | broadcast_inits.push_back(computation_->AddInstruction( |
| 3576 | HloInstruction::CreateBroadcast(reduce->shape().tuple_shapes(i), |
| 3577 | reduce->init_values()[i], {}))); |
| 3578 | } |
| 3579 | return ReplaceWithNewInstruction( |
| 3580 | reduce, HloInstruction::CreateTuple(broadcast_inits)); |
| 3581 | } else { |
| 3582 | return ReplaceWithNewInstruction( |
| 3583 | reduce, |
| 3584 | HloInstruction::CreateBroadcast(reduce_result_shape, init_value, {})); |
| 3585 | } |
| 3586 | } |
| 3587 | |
| 3588 | if (options_.is_layout_sensitive()) { |
| 3589 | return Status::OK(); |
| 3590 | } |
| 3591 | |
| 3592 | // If the reduction results in the same number of elements, then the only |
| 3593 | // possible side effect would be a reshape. Since the init_value is an |
| 3594 | // identity of the reduction function, we can therefore replace the reduce |
| 3595 | // with a simple reshape, ignoring the reduction function completely. |
| 3596 | if (ShapeUtil::ElementsIn(reduce_result_shape) == |
| 3597 | ShapeUtil::ElementsIn(arg->shape())) { |
| 3598 | if (multi_output_reduce) { |
| 3599 | std::vector<HloInstruction*> reshaped_args; |
| 3600 | int64 inputs = reduce->input_count(); |
| 3601 | for (int64 i = 0; i < inputs; ++i) { |
| 3602 | reshaped_args.push_back( |
| 3603 | computation_->AddInstruction(HloInstruction::CreateReshape( |
| 3604 | reduce->shape().tuple_shapes(i), reduce->inputs()[i]))); |
| 3605 | } |
| 3606 | return ReplaceWithNewInstruction( |
| 3607 | reduce, HloInstruction::CreateTuple(reshaped_args)); |
| 3608 | } else { |
| 3609 | return ReplaceWithNewInstruction( |
| 3610 | reduce, HloInstruction::CreateReshape(reduce_result_shape, arg)); |
| 3611 | } |
| 3612 | } |
| 3613 |
nothing calls this directly
no test coverage detected