| 128 | } |
| 129 | |
| 130 | StatusOr<HloInstruction*> RngExpander::ExpandInstruction(HloInstruction* rng) { |
| 131 | VLOG(2) << "Expand rng instruction " << rng->ToString(); |
| 132 | PrimitiveType old_primitive_type = rng->shape().element_type(); |
| 133 | if (primitive_util::BitWidth(old_primitive_type) < 32) { |
| 134 | TF_ASSIGN_OR_RETURN(rng, ConvertSmallFpRngToF32Rng(rng)); |
| 135 | } |
| 136 | HloComputation*& rng_computation = expanded_rng_instructions_[std::make_tuple( |
| 137 | rng->random_distribution(), rng->shape(), rng->operand(0)->shape(), |
| 138 | rng->operand(1)->shape())]; |
| 139 | if (!rng_computation) { |
| 140 | TF_ASSIGN_OR_RETURN(rng_computation, GetComputationForRng(rng)); |
| 141 | } |
| 142 | HloComputation* computation = rng->parent(); |
| 143 | |
| 144 | // A random number generated by the per module random number generator. |
| 145 | int64 module_random_value = rng->GetModule()->RandomNew64(); |
| 146 | |
| 147 | // A value specified by the configuration or generated by a global random |
| 148 | // number generator. |
| 149 | int64 module_config_seed = rng->parent()->parent()->config().seed(); |
| 150 | int64 global_random_value = |
| 151 | module_config_seed != 0 ? module_config_seed : GlobalRandomValue(); |
| 152 | |
| 153 | // Construct the key using the two random values above. |
| 154 | HloInstruction* key = MakeR0ConstantHlo<uint64>( |
| 155 | computation, module_random_value ^ global_random_value); |
| 156 | |
| 157 | const Shape u128_shape = ShapeUtil::MakeShape(xla::U64, {2}); |
| 158 | HloInstruction* state = |
| 159 | computation->AddInstruction(HloInstruction::CreateRngGetAndUpdateState( |
| 160 | u128_shape, GetNumberOf32bitUnits(rng->shape()))); |
| 161 | |
| 162 | VLOG(2) << "Rng key " << key->ToString(); |
| 163 | VLOG(2) << "Rng state " << state->ToString(); |
| 164 | |
| 165 | HloInstruction* new_rng = |
| 166 | computation->AddInstruction(HloInstruction::CreateCall( |
| 167 | rng->shape(), |
| 168 | {key, state, rng->mutable_operand(0), rng->mutable_operand(1)}, |
| 169 | rng_computation)); |
| 170 | |
| 171 | TF_RETURN_IF_ERROR(new_rng->CopyAllControlDepsFrom(rng)); |
| 172 | |
| 173 | TF_RETURN_IF_ERROR(rng->ReplaceAllUsesWith(new_rng)); |
| 174 | TF_RETURN_IF_ERROR(rng->DropAllControlDeps()); |
| 175 | |
| 176 | // Since rng is a side-effecting instruction, we can't rely on DCE to remove |
| 177 | // it. |
| 178 | TF_RETURN_IF_ERROR(computation->RemoveInstruction(rng)); |
| 179 | |
| 180 | // Returns nullptr to OpExpanderPass::Run to indicate the old rng instruction |
| 181 | // has been replaced with the new rng instruction. |
| 182 | return nullptr; |
| 183 | } |
| 184 | |
| 185 | } // namespace xla |
nothing calls this directly
no test coverage detected