Home
last modified time | relevance | path

Searched refs:while_hlo (Results 1 – 9 of 9) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dwhile_transformer_test.cc102 auto while_hlo = builder.AddInstruction( in BuildWhileInstruction() local
106 return while_hlo; in BuildWhileInstruction()
130 auto while_hlo = BuildWhileInstruction(condition, body, 0, 0); in TEST_F() local
131 auto result = ComputeWhileLoopTripCount(while_hlo); in TEST_F()
141 auto while_hlo = BuildWhileInstruction(condition, body, 1, 0); in TEST_F() local
142 auto result = ComputeWhileLoopTripCount(while_hlo); in TEST_F()
152 auto while_hlo = BuildWhileInstruction(condition, body, 0, 10); in TEST_F() local
153 auto result = ComputeWhileLoopTripCount(while_hlo); in TEST_F()
163 auto while_hlo = BuildWhileInstruction(condition, body, 0, 0); in TEST_F() local
164 auto result = ComputeWhileLoopTripCount(while_hlo); in TEST_F()
/external/tensorflow/tensorflow/compiler/xla/service/
Dcopy_insertion_test.cc650 auto while_hlo = builder.AddInstruction(HloInstruction::CreateWhile( in BuildWhileInstruction() local
653 return while_hlo; in BuildWhileInstruction()
658 auto while_hlo = builder.AddInstruction(HloInstruction::CreateWhile( in BuildWhileInstruction() local
661 return while_hlo; in BuildWhileInstruction()
766 auto while_hlo = builder->AddInstruction(HloInstruction::CreateWhile( in BuildWhileInstructionWithCustomInit() local
769 return while_hlo; in BuildWhileInstructionWithCustomInit()
798 auto while_hlo = BuildWhileInstruction(condition, body); in TEST_F() local
807 EXPECT_THAT(while_hlo->operand(0), in TEST_F()
826 auto while_hlo = BuildWhileInstruction(condition, body); in TEST_F() local
843 while_hlo->while_body()->root_instruction(), in TEST_F()
[all …]
Dbfloat16_propagation.cc164 HloInstruction* while_hlo) { in DetermineWhileComputationsPrecision() argument
165 CHECK_EQ(while_hlo->opcode(), HloOpcode::kWhile); in DetermineWhileComputationsPrecision()
172 HloComputation* body = while_hlo->while_body(); in DetermineWhileComputationsPrecision()
174 HloComputation* condition = while_hlo->while_condition(); in DetermineWhileComputationsPrecision()
177 body_root->shape(), [this, while_hlo, body_root]( in DetermineWhileComputationsPrecision()
182 if (OutputTypeAfterChange(while_hlo, index) == BF16) { in DetermineWhileComputationsPrecision()
187 << while_hlo->ToString(); in DetermineWhileComputationsPrecision()
Ddynamic_dimension_inference_test.cc521 HloInstruction* while_hlo = nullptr; in TEST_F() local
525 while_hlo = inst; in TEST_F()
528 ASSERT_NE(while_hlo, nullptr); in TEST_F()
531 EXPECT_EQ(while_hlo->shape().tuple_shapes_size(), 3); in TEST_F()
533 for (HloInstruction* inst : while_hlo->while_body()->instructions()) { in TEST_F()
540 EXPECT_EQ(inference_->GetDynamicSize(while_hlo, {0}, 0), size_param); in TEST_F()
541 EXPECT_EQ(inference_->GetDynamicSize(while_hlo, {1}, 0), size_param); in TEST_F()
Dbfloat16_propagation_test.cc599 auto while_hlo = builder.AddInstruction( in TEST_F() local
602 auto dot = builder.AddInstruction(CreateDot(shape, while_hlo, while_hlo)); in TEST_F()
665 auto while_hlo = builder.AddInstruction( in TEST_F() local
668 auto dot = builder.AddInstruction(CreateDot(shape, while_hlo, while_hlo)); in TEST_F()
742 auto while_hlo = builder.AddInstruction( in TEST_F() local
746 HloInstruction::CreateGetTupleElement(shape, while_hlo, 0)); in TEST_F()
748 HloInstruction::CreateGetTupleElement(shape, while_hlo, 1)); in TEST_F()
Dbfloat16_propagation.h111 void DetermineWhileComputationsPrecision(HloInstruction* while_hlo);
Dbuffer_assignment.cc1562 const HloInstruction* while_hlo = instruction; in BuildColocatedBufferSets() local
1564 while_hlo->shape(), in BuildColocatedBufferSets()
1565 [this, while_hlo, &points_to_analysis, buffer_size, in BuildColocatedBufferSets()
1570 AddBufferToColocatedSet(while_hlo->operand(0), index, in BuildColocatedBufferSets()
1573 AddBufferToColocatedSet(while_hlo, index, points_to_analysis, in BuildColocatedBufferSets()
1577 while_hlo->while_condition()->parameter_instruction(0), index, in BuildColocatedBufferSets()
1581 while_hlo->while_body()->parameter_instruction(0), index, in BuildColocatedBufferSets()
1585 while_hlo->while_body()->root_instruction(), index, in BuildColocatedBufferSets()
Dhlo_evaluator.cc1319 Status HloEvaluator::HandleWhile(HloInstruction* while_hlo) { in HandleWhile() argument
1320 HloComputation* cond_comp = while_hlo->while_condition(); in HandleWhile()
1321 HloComputation* body_comp = while_hlo->while_body(); in HandleWhile()
1323 auto lcv = GetEvaluatedLiteralFor(while_hlo->operand(0)).Clone(); in HandleWhile()
1334 while_hlo->name(), max_loop_iterations_); in HandleWhile()
1348 evaluated_[while_hlo] = std::move(lcv); in HandleWhile()
Dhlo_evaluator.h219 Status HandleWhile(HloInstruction* while_hlo) override;