Home
last modified time | relevance | path

Searched refs:branch_computation (Results 1 – 25 of 26) sorted by relevance

12

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dconditional_code_motion_test.cc284 const HloComputation* on_true = conditional->branch_computation(0); in TEST_F()
286 const HloComputation* on_false = conditional->branch_computation(1); in TEST_F()
352 const HloComputation* on_true = conditional->branch_computation(0); in TEST_F()
354 const HloComputation* on_false = conditional->branch_computation(1); in TEST_F()
421 const HloComputation* on_false = conditional->branch_computation(1); in TEST_F()
491 const HloComputation* on_true = conditional->branch_computation(0); in TEST_F()
496 const HloComputation* on_false = conditional->branch_computation(1); in TEST_F()
713 const HloComputation* on_true = conditional->branch_computation(0); in TEST_F()
715 const HloComputation* on_false = conditional->branch_computation(1); in TEST_F()
800 const HloComputation* on_true = conditional->branch_computation(0); in TEST_F()
[all …]
H A Dconditional_simplifier.cc125 if (conditional->branch_computation(branch) != computation) { in TryRemoveUnusedConditionalOperands()
146 conditional->branch_computation(branch) in TryRemoveUnusedConditionalOperands()
151 conditional->branch_computation(branch)->root_instruction()->shape())) in TryRemoveUnusedConditionalOperands()
152 << conditional->branch_computation(branch)->ToString(); in TryRemoveUnusedConditionalOperands()
167 auto branch_computation = in ReplaceRootWithEmptyTupleIfNoUsers() local
169 conditional_op->branch_computation(branch_id)->Clone()); in ReplaceRootWithEmptyTupleIfNoUsers()
170 conditional_op->set_branch_computation(branch_id, branch_computation); in ReplaceRootWithEmptyTupleIfNoUsers()
172 branch_computation->AddInstruction(HloInstruction::CreateTuple({})); in ReplaceRootWithEmptyTupleIfNoUsers()
173 branch_computation->set_root_instruction(new_empty_root, in ReplaceRootWithEmptyTupleIfNoUsers()
281 HloComputation* old_branch = conditional_op->branch_computation(branch_id); in RemoveUnusedTupleElements()
[all …]
H A Dconditional_code_motion.cc308 conditional->branch_computation(j)->AddInstruction( in CopyIntoConditional()
318 conditional->branch_computation(j)->root_instruction(); in CopyIntoConditional()
338 conditional->branch_computation(j)->AddInstruction( in CopyIntoConditional()
392 conditional->branch_computation(others)->root_instruction(); in FindSpecialConverts()
486 conditional->branch_computation(branch_num)->root_instruction(); in ConvertSpecialMove()
493 conditional->branch_computation(0)->root_instruction(); in ConvertSpecialMove()
522 old_root = conditional->branch_computation(branch)->root_instruction(); in ConvertSpecialMove()
561 HloComputation* cur_branch = conditional->branch_computation(branch); in ConvertSpecialMove()
683 auto computation = conditional->branch_computation(i); in MoveInstructionOut()
711 conditional->branch_computation(0)->root_instruction(); in MoveInstructionOut()
[all …]
H A Dhlo_ordering.cc100 a, a_ancestor->branch_computation(j))) { in GetExecutionConstraint()
104 b, a_ancestor->branch_computation(j))) { in GetExecutionConstraint()
204 b.defining_instruction()->branch_computation(j))) { in IsDefinedBefore()
347 conditional->branch_computation(j))) { in UsesBeforeValueDefinition()
H A Dlayout_assignment.cc738 computation_layouts_.find(instruction->branch_computation(0)) != in AddMandatoryConstraints()
747 instruction->branch_computation(j)->instruction_count(); in AddMandatoryConstraints()
759 instruction->branch_computation(largest_branch)) in AddMandatoryConstraints()
764 TF_RET_CHECK(instruction->branch_computation(j)->num_parameters() == 1); in AddMandatoryConstraints()
766 mutable_computation_constraints(instruction->branch_computation(k)) in AddMandatoryConstraints()
774 instruction->branch_computation(k), in AddMandatoryConstraints()
897 instruction->branch_computation(j)->root_instruction()->shape(), in CheckConditionalLayout()
900 << instruction->branch_computation(j)->root_instruction()->ToString(); in CheckConditionalLayout()
1071 auto branch_comp = instruction->branch_computation(operand_no - 1); in CopyOperandIfLayoutsDiffer()
1212 for (const auto branch_computation : branch_computations) { in CheckLayouts() local
[all …]
H A Dflatten_call_graph.cc54 CHECK_EQ(computation, instruction->branch_computation(b)); in ReplaceCalledComputation()
56 if (computation == instruction->branch_computation(b)) { in ReplaceCalledComputation()
H A Dbfloat16_propagation.cc208 auto branch = cond->branch_computation(i); in DetermineConditionalComputationsPrecision()
295 use.instruction->branch_computation(use.operand_number - 1) in AllUsersConsumeBF16()
501 adjust_computation(hlo->branch_computation(i), in AdjustCalledComputationParameters()
H A Dhlo_verifier.cc1337 conditional, conditional->branch_computation(j), 1)); in HandleConditional()
1339 conditional, j + 1, conditional->branch_computation(j), 0)); in HandleConditional()
1342 conditional->branch_computation(j)->root_instruction()->shape())); in HandleConditional()
2428 if (conditional->branch_computation(b)->num_parameters() != 1) { in HandleConditional()
2431 conditional->branch_computation(b)->name(), conditional->ToString(), in HandleConditional()
2432 conditional->branch_computation(b)->num_parameters()); in HandleConditional()
H A Ddynamic_dimension_inference.cc1481 HloComputation* branch_computation = hlo->branch_computation(branch_index); in HandleConditional() local
1483 HloComputation* new_computation = branch_computation; in HandleConditional()
1491 WidenComputation(branch_computation, new_operand->shape())); in HandleConditional()
H A Dar_crs_combiner_test.cc1678 auto branch_true = cond->branch_computation(0)->root_instruction(); in TEST_F()
1683 auto branch_false = cond->branch_computation(1)->root_instruction(); in TEST_F()
H A Dhlo_dataflow_analysis.cc780 conditional->branch_computation(j)->root_instruction()); in UpdateConditionalValueSet()
938 if (parameter->parent() == conditional->branch_computation(j)) { in UpdateParameterValueSet()
1261 user->branch_computation(j)->parameter_instruction(0)); in Propagate()
H A Dhlo_cost_analysis.cc1045 ProcessSubcomputation(conditional->branch_computation(0))); in HandleConditional()
1050 ProcessSubcomputation(conditional->branch_computation(j))); in HandleConditional()
H A Dcopy_insertion_test.cc3014 auto tuple6 = conditional18->branch_computation(1)->root_instruction(); in TEST_F()
3062 auto tuple6 = conditional18->branch_computation(1)->root_instruction(); in TEST_F()
3108 auto tuple6 = conditional18->branch_computation(1)->root_instruction(); in TEST_F()
3154 auto tuple6 = conditional18->branch_computation(1)->root_instruction(); in TEST_F()
H A Dar_crs_combiner.cc253 if (cond_instr->branch_computation(i)->parameter_instruction(0) == in GetAllTuples()
H A Dhlo_module_group_metadata.cc326 tracked_instructions_[hlo->branch_computation(b)] = in RecordInstructions()
H A Dhlo_module.cc207 replacements, instruction->branch_computation(b), nullptr); in ReplaceComputations()
H A Dhlo_instruction.cc2543 if (!eq_computations(branch_computation(j), in IdenticalSlowPath()
2544 other.branch_computation(j))) { in IdenticalSlowPath()
2883 HloComputation* HloInstruction::branch_computation(int b) const { in branch_computation() function in xla::HloInstruction
H A Dhlo_instruction.h1482 HloComputation* branch_computation(int b) const;
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/
H A Dir_emitter.cc2822 HloComputation* br_computation = conditional->branch_computation(b); in HandleConditional()
2851 EmitGlobalCall(*conditional->branch_computation(0), in HandleConditional()
2855 EmitGlobalCall(*conditional->branch_computation(1), in HandleConditional()
2897 EmitGlobalCall(*conditional->branch_computation(num_branches - 1), in HandleConditional()
2911 EmitGlobalCall(*conditional->branch_computation(b), in HandleConditional()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/
H A Dxla_builder_test.cc824 ->branch_computation(0) in TEST_F()
829 ->branch_computation(1) in TEST_F()
H A Dxla_builder.cc2556 for (const XlaComputation* branch_computation : branch_computations) { in ConditionalImpl() local
2557 AddCalledComputation(*branch_computation, &instr); in ConditionalImpl()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/
H A Dhlo_function_importer.cc936 {instruction->branch_computation(0)->root_instruction()}, &rets)); in ImportInstructionImpl()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dspmd_partitioner_test.cc716 HloInstruction* root_b0 = root->branch_computation(0)->root_instruction(); in TEST_F()
722 HloInstruction* root_b1 = root->branch_computation(1)->root_instruction(); in TEST_F()
3647 auto then_branch_root = root->branch_computation(0)->root_instruction(); in TEST_F()
3653 auto else_branch_root = root->branch_computation(1)->root_instruction(); in TEST_F()
H A Dspmd_partitioner.cc3641 HloComputation* computation = hlo->branch_computation(i); in HandleConditional()
3653 HloComputation* computation = hlo->branch_computation(i); in HandleConditional()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dir_emitter_unnested.cc649 mlir::Region* branch_computation = &conditional.getBranches()[j]; in EmitConditional() local
653 TF_RETURN_IF_ERROR(ir_emitter->EmitLmhloRegion(branch_computation)); in EmitConditional()

12