Home
last modified time | relevance | path

Searched refs:xla_while (Results 1 – 24 of 24) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_module_dce.cc47 const auto* xla_while = instruction; in RunWhileDCE() local
48 auto* while_body_comp = xla_while->while_body(); in RunWhileDCE()
52 if (!xla_while->shape().IsTuple() || in RunWhileDCE()
56 VLOG(1) << "WhileDCE SKIP while: " << xla_while->ToString(); in RunWhileDCE()
62 ShapeUtil::TupleElementCount(xla_while->shape()); in RunWhileDCE()
64 if (liveness->IsLive(xla_while, {i})) { in RunWhileDCE()
68 << " while: " << xla_while->name() << " tuple_index: " << i; in RunWhileDCE()
Dhlo_ordering_test.cc150 auto xla_while = builder.AddInstruction( in TEST_F() local
155 EXPECT_TRUE(ordering.ExecutesBefore(constant, xla_while)); in TEST_F()
163 EXPECT_FALSE(ordering.ExecutesBefore(xla_while, body_param)); in TEST_F()
164 EXPECT_FALSE(ordering.ExecutesBefore(xla_while, cond_param)); in TEST_F()
165 EXPECT_FALSE(ordering.ExecutesBefore(body_param, xla_while)); in TEST_F()
166 EXPECT_FALSE(ordering.ExecutesBefore(cond_param, xla_while)); in TEST_F()
233 auto xla_while = builder.AddInstruction( in TEST_F() local
236 scalar_shape, HloOpcode::kAdd, constant, xla_while)); in TEST_F()
246 dataflow->GetValueDefinedAt(xla_while))); in TEST_F()
249 dataflow->GetValueDefinedAt(xla_while), *dataflow)); in TEST_F()
[all …]
Dhlo_alias_analysis_test.cc365 auto xla_while = builder.AddInstruction( in TEST_F() local
368 HloInstruction::CreateGetTupleElement(scalar_shape_, xla_while, 0)); in TEST_F()
370 HloInstruction::CreateGetTupleElement(scalar_shape_, xla_while, 1)); in TEST_F()
390 GetValuesInBuffer(analysis.GetUniqueBufferAt(xla_while, /*index=*/{1})), in TEST_F()
392 GetValueDefinedAt(xla_while, /*index=*/{1}), in TEST_F()
399 analysis.GetUniqueBufferAt(xla_while, /*index=*/{1}).ComputePositions(), in TEST_F()
401 HloPosition{param, {1}}, HloPosition{xla_while, {1}}, in TEST_F()
556 auto xla_while = builder.AddInstruction( in TEST_F() local
565 analysis.GetUniqueBufferAt(xla_while, /*index=*/{}).ComputePositions(), in TEST_F()
566 UnorderedElementsAre(HloPosition{tuple, {}}, HloPosition{xla_while, {}}, in TEST_F()
[all …]
Dhlo_schedule_test.cc252 const HloInstruction* xla_while = in TEST_F() local
254 HloComputation* body = xla_while->while_body(); in TEST_F()
255 HloComputation* cond = xla_while->while_condition(); in TEST_F()
321 HloInstruction* xla_while = in TEST_F() local
323 HloInstruction* init = xla_while->mutable_operand(0); in TEST_F()
327 TF_ASSERT_OK(xla_while->ReplaceAllUsesWith(init)); in TEST_F()
Dcopy_insertion.cc169 const HloInstruction* xla_while, in IndicesToCopyForWhile() argument
171 DCHECK(ShapeUtil::Compatible(indices_to_copy->shape(), xla_while->shape())); in IndicesToCopyForWhile()
174 const HloInstruction* init = xla_while->operand(0); in IndicesToCopyForWhile()
180 dataflow.GetValueSet(xla_while, index).values().size() > 1) { in IndicesToCopyForWhile()
186 should_copy = dataflow.GetUniqueValueAt(xla_while, index) != in IndicesToCopyForWhile()
257 HloInstruction* xla_while) { in AddCopiesForWhile() argument
258 VLOG(2) << "Adding copies for kWhile instruction " << xla_while->name(); in AddCopiesForWhile()
259 TF_RET_CHECK(xla_while->opcode() == HloOpcode::kWhile); in AddCopiesForWhile()
261 ShapeTree<bool> indices_to_copy(xla_while->shape()); in AddCopiesForWhile()
262 if (!IndicesToCopyForWhile(alias_analysis.dataflow_analysis(), xla_while, in AddCopiesForWhile()
[all …]
Dhlo_ordering.cc195 const HloInstruction* xla_while = use.instruction; in UseIsBeforeValueDefinition() local
197 xla_while->while_body()) || in UseIsBeforeValueDefinition()
199 xla_while->while_condition())) { in UseIsBeforeValueDefinition()
210 const HloInstruction* xla_while = value.defining_instruction(); in UseIsBeforeValueDefinition() local
212 xla_while->while_body()) || in UseIsBeforeValueDefinition()
214 xla_while->while_condition())) { in UseIsBeforeValueDefinition()
Dcopy_insertion_test.cc732 auto xla_while = BuildWhileInstructionWithCustomInit(loop_state_shape_, in BuildWhileInstruction_InitPointsToInterfering() local
737 auto gte = xla_while->parent()->AddInstruction( in BuildWhileInstruction_InitPointsToInterfering()
739 ShapeUtil::GetSubshape(xla_while->shape(), {1}), xla_while, 1)); in BuildWhileInstruction_InitPointsToInterfering()
740 auto sub = xla_while->parent()->AddInstruction(HloInstruction::CreateBinary( in BuildWhileInstruction_InitPointsToInterfering()
742 auto gte0 = xla_while->parent()->AddInstruction( in BuildWhileInstruction_InitPointsToInterfering()
744 ShapeUtil::GetSubshape(xla_while->shape(), {0}), xla_while, 0)); in BuildWhileInstruction_InitPointsToInterfering()
745 auto tuple = xla_while->parent()->AddInstruction( in BuildWhileInstruction_InitPointsToInterfering()
748 xla_while->parent()->set_root_instruction(tuple); in BuildWhileInstruction_InitPointsToInterfering()
750 return xla_while; in BuildWhileInstruction_InitPointsToInterfering()
1331 auto xla_while = builder.AddInstruction( in TEST_F() local
[all …]
Dhlo_verifier.cc646 Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) { in HandleWhile() argument
648 CheckParameterCount(xla_while, xla_while->while_body(), 1)); in HandleWhile()
650 CheckParameterCount(xla_while, xla_while->while_condition(), 1)); in HandleWhile()
652 CheckOperandAndParameter(xla_while, 0, xla_while->while_body(), 0)); in HandleWhile()
654 CheckOperandAndParameter(xla_while, 0, xla_while->while_condition(), 0)); in HandleWhile()
656 xla_while->while_condition()->root_instruction()->shape(); in HandleWhile()
665 return CheckShape(xla_while, in HandleWhile()
666 xla_while->while_body()->root_instruction()->shape()); in HandleWhile()
1353 Status HandleWhile(HloInstruction* xla_while) override { in HandleWhile() argument
1354 auto* while_cond = xla_while->while_condition(); in HandleWhile()
[all …]
Dhlo_dataflow_analysis_test.cc515 auto xla_while = builder.AddInstruction( in TEST_P() local
528 EXPECT_FALSE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{0})); in TEST_P()
533 EXPECT_TRUE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{1})); in TEST_P()
534 EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{1}).is_phi()); in TEST_P()
543 HloUse{xla_while, 0, {0}})); in TEST_P()
547 EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{1}) in TEST_P()
554 EXPECT_FALSE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{0})); in TEST_P()
555 EXPECT_FALSE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{1})); in TEST_P()
805 auto xla_while = builder.AddInstruction( in TEST_P() local
820 EXPECT_TRUE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{0})); in TEST_P()
[all …]
Dhlo_liveness_analysis.cc204 auto* xla_while = callsite.instruction(); in PropagateLivenessToParameterCallers() local
209 MarkLiveAtIndex(xla_while, shape_index, live_index_map, worklist, in PropagateLivenessToParameterCallers()
212 MarkLiveAtIndex(xla_while->while_body()->root_instruction(), in PropagateLivenessToParameterCallers()
215 MarkLiveAtIndex(xla_while->operand(0), shape_index, live_index_map, in PropagateLivenessToParameterCallers()
Dhlo_dataflow_analysis.cc586 bool HloDataflowAnalysis::UpdateWhileValueSet(HloInstruction* xla_while) { in UpdateWhileValueSet() argument
587 CHECK_EQ(xla_while->opcode(), HloOpcode::kWhile); in UpdateWhileValueSet()
589 &GetInstructionValueSet(xla_while->while_body()->root_instruction()), in UpdateWhileValueSet()
590 &GetInstructionValueSet(xla_while->operand(0))}; in UpdateWhileValueSet()
592 return Phi(xla_while, inputs); in UpdateWhileValueSet()
594 return GetInstructionValueSet(xla_while).AssignUnionOf(inputs); in UpdateWhileValueSet()
Ddfs_hlo_visitor_with_default.h188 Status HandleWhile(HloInstructionPtr xla_while) override { in HandleWhile() argument
189 return DefaultAction(xla_while); in HandleWhile()
Dhlo_cost_analysis.cc663 Status HloCostAnalysis::HandleWhile(const HloInstruction* xla_while) { in HandleWhile() argument
668 ProcessUnnestedSubcomputation(xla_while->while_body())); in HandleWhile()
672 ProcessUnnestedSubcomputation(xla_while->while_condition())); in HandleWhile()
Dhlo_dataflow_analysis.h195 bool UpdateWhileValueSet(HloInstruction* xla_while);
Dhlo_cost_analysis.h110 Status HandleWhile(const HloInstruction* xla_while) override;
Dhlo_verifier.h86 Status HandleWhile(HloInstruction* xla_while) override;
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_copy_insertion_test.cc88 auto xla_while = builder.AddInstruction( in TEST_F() local
97 EXPECT_THAT(xla_while->operand(0), op::Copy(op::Parameter())); in TEST_F()
Dir_emitter.cc2296 Status IrEmitter::HandleWhile(HloInstruction* xla_while) { in HandleWhile() argument
2298 HloComputation* condition = xla_while->while_condition(); in HandleWhile()
2305 xla_while->shape(), in HandleWhile()
2306 [this, &xla_while](const Shape& /*subshape*/, in HandleWhile()
2323 TF_RETURN_IF_ERROR(check(xla_while, xla_while->operand(0), index)); in HandleWhile()
2325 xla_while, xla_while->while_condition()->parameter_instruction(0), in HandleWhile()
2328 check(xla_while, xla_while->while_body()->parameter_instruction(0), in HandleWhile()
2331 xla_while, xla_while->while_body()->root_instruction(), index)); in HandleWhile()
2336 const HloInstruction* init = xla_while->operand(0); in HandleWhile()
2337 emitted_value_[xla_while] = GetEmittedValueFor(init); in HandleWhile()
[all …]
Dcpu_compiler.cc226 Status HandleWhile(HloInstruction* xla_while) override { in HandleWhile() argument
227 TF_RETURN_IF_ERROR(DefaultAction(xla_while)); in HandleWhile()
232 xla_while->while_condition()->Accept(&candidates_for_condition)); in HandleWhile()
236 TF_RETURN_IF_ERROR(xla_while->while_body()->Accept(&candidates_for_body)); in HandleWhile()
Dir_emitter.h178 Status HandleWhile(HloInstruction* xla_while) override;
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter_unnested.cc977 Status IrEmitterUnnested::HandleWhile(HloInstruction* xla_while) { in HandleWhile() argument
978 HloComputation* condition = xla_while->while_condition(); in HandleWhile()
983 auto config = xla_while->backend_config<WhileLoopBackendConfig>(); in HandleWhile()
986 BuildForThunk(xla_while, config.ValueOrDie().known_trip_count().n())); in HandleWhile()
988 AddThunkToThunkSequence(BuildWhileThunk(xla_while)); in HandleWhile()
2055 const HloInstruction* xla_while, in CheckWhileBuffersShareAllocation() argument
2058 xla_while->shape(), in CheckWhileBuffersShareAllocation()
2061 xla_while->while_condition()->parameter_instruction(0); in CheckWhileBuffersShareAllocation()
2062 const HloComputation* body = xla_while->while_body(); in CheckWhileBuffersShareAllocation()
2066 xla_while, xla_while->operand(0), index, buffer_assignment)); in CheckWhileBuffersShareAllocation()
[all …]
Dir_emitter_unnested.h172 Status HandleWhile(HloInstruction* xla_while) override;
/external/tensorflow/tensorflow/compiler/xla/tests/
Dwhile_test.cc449 auto xla_while = While(condition, body, init); in XLA_TEST_F() local
452 Add(GetTupleElement(xla_while, 1), GetTupleElement(xla_while, 2)); in XLA_TEST_F()
453 auto result = Add(add12, GetTupleElement(xla_while, 3)); in XLA_TEST_F()
/external/tensorflow/tensorflow/compiler/tf2xla/python/
Dxla.py420 while_loop = gen_xla_ops.xla_while