/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
H A D | conditional_code_motion_test.cc | 284 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 D | conditional_simplifier.cc | 125 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 D | conditional_code_motion.cc | 308 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 D | hlo_ordering.cc | 100 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 D | layout_assignment.cc | 738 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 D | flatten_call_graph.cc | 54 CHECK_EQ(computation, instruction->branch_computation(b)); in ReplaceCalledComputation() 56 if (computation == instruction->branch_computation(b)) { in ReplaceCalledComputation()
|
H A D | bfloat16_propagation.cc | 208 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 D | hlo_verifier.cc | 1337 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 D | dynamic_dimension_inference.cc | 1481 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 D | ar_crs_combiner_test.cc | 1678 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 D | hlo_dataflow_analysis.cc | 780 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 D | hlo_cost_analysis.cc | 1045 ProcessSubcomputation(conditional->branch_computation(0))); in HandleConditional() 1050 ProcessSubcomputation(conditional->branch_computation(j))); in HandleConditional()
|
H A D | copy_insertion_test.cc | 3014 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 D | ar_crs_combiner.cc | 253 if (cond_instr->branch_computation(i)->parameter_instruction(0) == in GetAllTuples()
|
H A D | hlo_module_group_metadata.cc | 326 tracked_instructions_[hlo->branch_computation(b)] = in RecordInstructions()
|
H A D | hlo_module.cc | 207 replacements, instruction->branch_computation(b), nullptr); in ReplaceComputations()
|
H A D | hlo_instruction.cc | 2543 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 D | hlo_instruction.h | 1482 HloComputation* branch_computation(int b) const;
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
H A D | ir_emitter.cc | 2822 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 D | xla_builder_test.cc | 824 ->branch_computation(0) in TEST_F() 829 ->branch_computation(1) in TEST_F()
|
H A D | xla_builder.cc | 2556 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 D | hlo_function_importer.cc | 936 {instruction->branch_computation(0)->root_instruction()}, &rets)); in ImportInstructionImpl()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
H A D | spmd_partitioner_test.cc | 716 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 D | spmd_partitioner.cc | 3641 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 D | ir_emitter_unnested.cc | 649 mlir::Region* branch_computation = &conditional.getBranches()[j]; in EmitConditional() local 653 TF_RETURN_IF_ERROR(ir_emitter->EmitLmhloRegion(branch_computation)); in EmitConditional()
|