| 98 | } |
| 99 | |
| 100 | Status SoftmaxExpanderVisitor::HandleSoftmax( |
| 101 | HloInstruction* softmax) { |
| 102 | std::vector<HloInstruction*> added_instructions; |
| 103 | auto add = [&](std::unique_ptr<HloInstruction> inst) { |
| 104 | HloInstruction* added_inst = computation_->AddInstruction(std::move(inst)); |
| 105 | added_inst->set_metadata(softmax->metadata()); |
| 106 | added_instructions.push_back(added_inst); |
| 107 | return added_inst; |
| 108 | }; |
| 109 | auto add_binary = [&](const Shape& shape, const HloOpcode opcode, |
| 110 | HloInstruction* a, HloInstruction* b) { |
| 111 | return add(HloInstruction::CreateBinary(shape, opcode, a, b)); |
| 112 | }; |
| 113 | int64 instruction_count_before = computation_->instruction_count(); |
| 114 | |
| 115 | // Expand batch norm training into smaller HLO ops. |
| 116 | HloInstruction* operand = softmax->mutable_operand(0); |
| 117 | const Shape operand_shape = operand->shape(); |
| 118 | PrimitiveType ptype = operand_shape.element_type(); |
| 119 | int64 feature_index = softmax->softmax_feature_index(); |
| 120 | std::vector<int64> dimensions_without_feature; |
| 121 | std::vector<int64> reduction_dimensions; |
| 122 | |
| 123 | for (int64 i = 0; i < operand_shape.rank(); ++i) { |
| 124 | if (i != feature_index) { |
| 125 | dimensions_without_feature.push_back(i); |
| 126 | reduction_dimensions.push_back(operand_shape.dimensions(i)); |
| 127 | } |
| 128 | } |
| 129 | const Shape reduce_shape = ShapeUtil::MakeShape(ptype, reduction_dimensions); |
| 130 | |
| 131 | // max(operand) |
| 132 | auto minf_literal = LiteralUtil::MinValue(ptype); |
| 133 | auto minf = add(HloInstruction::CreateConstant(std::move(minf_literal))); |
| 134 | HloComputation* max_reduce_computation = |
| 135 | GetOrCreateScalarMaxComputation(ptype); |
| 136 | auto max = add(HloInstruction::CreateReduce(reduce_shape, operand, minf, |
| 137 | {feature_index}, |
| 138 | max_reduce_computation)); |
| 139 | auto max_broadcasted = |
| 140 | add(HloInstruction::CreateBroadcast(operand_shape, max, |
| 141 | dimensions_without_feature)); |
| 142 | // operand - max(operand) |
| 143 | auto sub = |
| 144 | add_binary(operand_shape, HloOpcode::kSubtract, operand, max_broadcasted); |
| 145 | // e^(operand - max(operand)) |
| 146 | auto exp = add(HloInstruction::CreateUnary(operand_shape, HloOpcode::kExp, |
| 147 | sub)); |
| 148 | // sum(e^(operand - max(operand))) |
| 149 | auto zero_literal = LiteralUtil::CreateR0(0.0f); |
| 150 | TF_ASSIGN_OR_RETURN(zero_literal, zero_literal.Convert(ptype)); |
| 151 | auto zero = add(HloInstruction::CreateConstant(std::move(zero_literal))); |
| 152 | HloComputation* add_reduce_computation = |
| 153 | GetOrCreateScalarAddComputation(ptype); |
| 154 | auto sum = add(HloInstruction::CreateReduce(reduce_shape, exp, zero, |
| 155 | {feature_index}, |
| 156 | add_reduce_computation)); |
| 157 | auto sum_broadcasted = |
nothing calls this directly
no test coverage detected