Searched refs:while_body_root (Results 1 – 5 of 5) sorted by relevance
55 HloInstruction* while_body_root = while_body->root_instruction(); in RemoveDeadTupleIndices() local92 if (user == while_body_root) { in RemoveDeadTupleIndices()110 user->users().front() == while_body_root) in RemoveDeadTupleIndices()138 while_body_root->mutable_operand(old_idx)); in RemoveDeadTupleIndices()141 while_body_root, HloInstruction::CreateTuple(new_while_body_root_elems)); in RemoveDeadTupleIndices()227 HloInstruction* while_body_root = while_body->root_instruction(); in TryRemoveDeadWhileParams() local234 if (while_body_root->opcode() != HloOpcode::kTuple) { in TryRemoveDeadWhileParams()278 if (user->user_count() == 1 && user->users().front() == while_body_root && in TryRemoveDeadWhileParams()279 while_body_root->operand_index(user) == user->tuple_index() && in TryRemoveDeadWhileParams()280 absl::c_count(while_body_root->operands(), user) == 1) { in TryRemoveDeadWhileParams()[all …]
135 auto* while_body_root = while_body->root_instruction(); in GetAuxiliaryLoopInductionVars() local136 if (while_body_root->opcode() != HloOpcode::kTuple) { in GetAuxiliaryLoopInductionVars()137 VLOG(2) << "While body root is not a tuple:" << while_body_root->ToString(); in GetAuxiliaryLoopInductionVars()142 for (const HloInstruction* operand : while_body_root->operands()) { in GetAuxiliaryLoopInductionVars()271 auto* while_body_root = while_body->root_instruction(); in GetLoopInductionVarTupleIdx() local272 if (while_body_root->opcode() != HloOpcode::kTuple) { in GetLoopInductionVarTupleIdx()274 << while_body_root->ToString(); in GetLoopInductionVarTupleIdx()278 auto* while_body_inc = while_body_root->operand(*indvar_tuple_idx); in GetLoopInductionVarTupleIdx()591 auto* while_body_root = while_body->root_instruction(); in ComputeWhileLoopTripCountUpperBound() local592 if (while_body_root->opcode() != HloOpcode::kTuple) { in ComputeWhileLoopTripCountUpperBound()[all …]
53 auto* while_body_root = while_body_comp->root_instruction(); in RunWhileDCE() local56 while_body_root->opcode() != HloOpcode::kTuple) { in RunWhileDCE()84 while_body_root->ReplaceOperandWith(i, pass_thru_gte)); in RunWhileDCE()
30 HloInstruction* while_body_root, int64 tuple_index) { in ReplaceUsesWhileKeepingLoopInvariance() argument31 CHECK_EQ(while_body_root->opcode(), HloOpcode::kTuple); in ReplaceUsesWhileKeepingLoopInvariance()40 !(user == while_body_root && i == tuple_index)) { in ReplaceUsesWhileKeepingLoopInvariance()
54 auto* while_body_root = while_body_comp->root_instruction(); in WhileBodyHasPassThroughTupleElement() local55 if (while_body_root->opcode() != HloOpcode::kTuple) { in WhileBodyHasPassThroughTupleElement()58 auto* operand = while_body_root->operand(tuple_index); in WhileBodyHasPassThroughTupleElement()