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.cc50 const auto* xla_while = instruction; in RunWhileDCE() local
51 auto* while_body_comp = xla_while->while_body(); in RunWhileDCE()
55 if (!xla_while->shape().IsTuple() || in RunWhileDCE()
59 VLOG(1) << "WhileDCE SKIP while: " << xla_while->ToString(); in RunWhileDCE()
65 ShapeUtil::TupleElementCount(xla_while->shape()); in RunWhileDCE()
68 if (liveness->IsLive(xla_while, {i})) { in RunWhileDCE()
72 << " while: " << xla_while->name() << " tuple_index: " << i; in RunWhileDCE()
Dhlo_ordering_test.cc149 auto xla_while = builder.AddInstruction( in TEST_F() local
154 EXPECT_TRUE(ordering.ExecutesBefore(constant, xla_while)); in TEST_F()
162 EXPECT_FALSE(ordering.ExecutesBefore(xla_while, body_param)); in TEST_F()
163 EXPECT_FALSE(ordering.ExecutesBefore(xla_while, cond_param)); in TEST_F()
164 EXPECT_FALSE(ordering.ExecutesBefore(body_param, xla_while)); in TEST_F()
165 EXPECT_FALSE(ordering.ExecutesBefore(cond_param, xla_while)); in TEST_F()
232 auto xla_while = builder.AddInstruction( in TEST_F() local
235 scalar_shape, HloOpcode::kAdd, constant, xla_while)); in TEST_F()
245 dataflow->GetValueDefinedAt(xla_while))); in TEST_F()
248 dataflow->GetValueDefinedAt(xla_while), *dataflow)); in TEST_F()
[all …]
Dhlo_alias_analysis_test.cc359 auto xla_while = builder.AddInstruction( in TEST_F() local
362 HloInstruction::CreateGetTupleElement(scalar_shape_, xla_while, 0)); in TEST_F()
364 HloInstruction::CreateGetTupleElement(scalar_shape_, xla_while, 1)); in TEST_F()
382 GetValuesInBuffer(analysis.GetUniqueBufferAt(xla_while, /*index=*/{1})), in TEST_F()
384 GetValueDefinedAt(xla_while, /*index=*/{1}), in TEST_F()
391 analysis.GetUniqueBufferAt(xla_while, /*index=*/{1}).ComputePositions(), in TEST_F()
393 HloPosition{param, {1}}, HloPosition{xla_while, {1}}, in TEST_F()
548 auto xla_while = builder.AddInstruction( in TEST_F() local
557 analysis.GetUniqueBufferAt(xla_while, /*index=*/{}).ComputePositions(), in TEST_F()
558 UnorderedElementsAre(HloPosition{tuple, {}}, HloPosition{xla_while, {}}, in TEST_F()
[all …]
Dhlo_schedule_test.cc251 const HloInstruction* xla_while = in TEST_F() local
253 HloComputation* body = xla_while->while_body(); in TEST_F()
254 HloComputation* cond = xla_while->while_condition(); in TEST_F()
320 HloInstruction* xla_while = in TEST_F() local
322 HloInstruction* init = xla_while->mutable_operand(0); in TEST_F()
326 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()
280 HloInstruction* xla_while) { in AddCopiesForWhile() argument
281 VLOG(2) << "Adding copies for kWhile instruction " << xla_while->name(); in AddCopiesForWhile()
282 TF_RET_CHECK(xla_while->opcode() == HloOpcode::kWhile); in AddCopiesForWhile()
284 ShapeTree<bool> indices_to_copy(xla_while->shape()); in AddCopiesForWhile()
285 if (!IndicesToCopyForWhile(alias_analysis.dataflow_analysis(), xla_while, in AddCopiesForWhile()
[all …]
Dhlo_ordering.cc272 const HloInstruction* xla_while = use.instruction; in UsesBeforeValueDefinition() local
274 xla_while->while_body())) { in UsesBeforeValueDefinition()
280 xla_while->while_condition())) { in UsesBeforeValueDefinition()
282 xla_while->while_condition()->parameter_instruction(0)) { in UsesBeforeValueDefinition()
297 const HloInstruction* xla_while = value.defining_instruction(); in UsesBeforeValueDefinition() local
299 xla_while->while_body()) || in UsesBeforeValueDefinition()
301 xla_while->while_condition())) { in UsesBeforeValueDefinition()
Dloop_schedule_linearizer.cc71 HloInstruction* xla_while, HloAliasAnalysis& alias_analysis) { in AddControlEdgesForLoopWrites() argument
73 HloComputation* body = xla_while->while_body(); in AddControlEdgesForLoopWrites()
83 ShapeTree<bool> indices_to_copy(xla_while->shape()); in AddControlEdgesForLoopWrites()
Dcopy_insertion_test.cc738 auto xla_while = BuildWhileInstructionWithCustomInit(loop_state_shape_, in BuildWhileInstruction_InitPointsToInterfering() local
743 auto gte = xla_while->parent()->AddInstruction( in BuildWhileInstruction_InitPointsToInterfering()
745 ShapeUtil::GetSubshape(xla_while->shape(), {1}), xla_while, 1)); in BuildWhileInstruction_InitPointsToInterfering()
746 auto sub = xla_while->parent()->AddInstruction(HloInstruction::CreateBinary( in BuildWhileInstruction_InitPointsToInterfering()
748 auto gte0 = xla_while->parent()->AddInstruction( in BuildWhileInstruction_InitPointsToInterfering()
750 ShapeUtil::GetSubshape(xla_while->shape(), {0}), xla_while, 0)); in BuildWhileInstruction_InitPointsToInterfering()
751 auto tuple = xla_while->parent()->AddInstruction( in BuildWhileInstruction_InitPointsToInterfering()
754 xla_while->parent()->set_root_instruction(tuple); in BuildWhileInstruction_InitPointsToInterfering()
756 return xla_while; in BuildWhileInstruction_InitPointsToInterfering()
1541 auto xla_while = builder.AddInstruction( in TEST_F() local
[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_verifier.cc920 Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) { in HandleWhile() argument
922 CheckParameterCount(xla_while, xla_while->while_body(), 1)); in HandleWhile()
924 CheckParameterCount(xla_while, xla_while->while_condition(), 1)); in HandleWhile()
926 CheckOperandAndParameter(xla_while, 0, xla_while->while_body(), 0)); in HandleWhile()
928 CheckOperandAndParameter(xla_while, 0, xla_while->while_condition(), 0)); in HandleWhile()
930 xla_while->while_condition()->root_instruction()->shape(); in HandleWhile()
940 return CheckShape(xla_while, in HandleWhile()
941 xla_while->while_body()->root_instruction()->shape()); in HandleWhile()
1815 Status HandleWhile(HloInstruction* xla_while) override { in HandleWhile() argument
1816 auto* while_cond = xla_while->while_condition(); in HandleWhile()
[all …]
Dhlo_dataflow_analysis_test.cc422 auto xla_while = builder.AddInstruction( in TEST_P() local
435 EXPECT_FALSE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{0})); in TEST_P()
440 EXPECT_TRUE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{1})); in TEST_P()
441 EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{1}).is_phi()); in TEST_P()
450 HloUse{xla_while, 0, {0}})); in TEST_P()
454 EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{1}) in TEST_P()
461 EXPECT_FALSE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{0})); in TEST_P()
462 EXPECT_FALSE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{1})); in TEST_P()
858 auto xla_while = builder.AddInstruction( in TEST_P() local
873 EXPECT_TRUE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{0})); in TEST_P()
[all …]
Ddfs_hlo_visitor_with_default.h210 Status HandleWhile(HloInstructionPtr xla_while) override { in HandleWhile() argument
211 return DefaultAction(xla_while); in HandleWhile()
Dhlo_dataflow_analysis.cc724 bool HloDataflowAnalysis::UpdateWhileValueSet(HloInstruction* xla_while) { in UpdateWhileValueSet() argument
725 CHECK_EQ(xla_while->opcode(), HloOpcode::kWhile); in UpdateWhileValueSet()
727 &GetInstructionValueSet(xla_while->while_body()->root_instruction()), in UpdateWhileValueSet()
728 &GetInstructionValueSet(xla_while->operand(0))}; in UpdateWhileValueSet()
730 return Phi(xla_while, inputs); in UpdateWhileValueSet()
732 return GetInstructionValueSet(xla_while).AssignUnionOf(inputs); in UpdateWhileValueSet()
Dhlo_dataflow_analysis.h230 bool UpdateWhileValueSet(HloInstruction* xla_while);
Dhlo_cost_analysis.h120 Status HandleWhile(const HloInstruction* xla_while) override;
Dhlo_cost_analysis.cc960 Status HloCostAnalysis::HandleWhile(const HloInstruction* xla_while) { in HandleWhile() argument
965 ProcessSubcomputation(xla_while->while_body())); in HandleWhile()
968 ProcessSubcomputation(xla_while->while_condition())); in HandleWhile()
Dhlo_verifier.h95 Status HandleWhile(HloInstruction* xla_while) override;
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dir_emitter.cc2357 Status IrEmitter::HandleWhile(HloInstruction* xla_while) { in HandleWhile() argument
2359 HloComputation* condition = xla_while->while_condition(); in HandleWhile()
2366 xla_while->shape(), in HandleWhile()
2367 [this, &xla_while](const Shape& /*subshape*/, in HandleWhile()
2384 TF_RETURN_IF_ERROR(check(xla_while, xla_while->operand(0), index)); in HandleWhile()
2386 xla_while, xla_while->while_condition()->parameter_instruction(0), in HandleWhile()
2389 check(xla_while, xla_while->while_body()->parameter_instruction(0), in HandleWhile()
2392 xla_while, xla_while->while_body()->root_instruction(), index)); in HandleWhile()
2397 const HloInstruction* init = xla_while->operand(0); in HandleWhile()
2398 emitted_value_[xla_while] = GetEmittedValueFor(init); in HandleWhile()
[all …]
Dcpu_compiler.cc262 Status HandleWhile(HloInstruction* xla_while) override { in HandleWhile() argument
263 TF_RETURN_IF_ERROR(DefaultAction(xla_while)); in HandleWhile()
268 xla_while->while_condition()->Accept(&candidates_for_condition)); in HandleWhile()
272 TF_RETURN_IF_ERROR(xla_while->while_body()->Accept(&candidates_for_body)); in HandleWhile()
Dir_emitter.h164 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/xla/service/gpu/
Dir_emitter_unnested.cc2431 Status IrEmitterUnnested::HandleWhile(HloInstruction* xla_while) { in HandleWhile() argument
2432 HloComputation* condition = xla_while->while_condition(); in HandleWhile()
2437 auto config = xla_while->backend_config<WhileLoopBackendConfig>(); in HandleWhile()
2441 BuildForThunk(xla_while, config.ValueOrDie().known_trip_count().n())); in HandleWhile()
2444 TF_ASSIGN_OR_RETURN(auto thunk, BuildWhileThunk(xla_while)); in HandleWhile()
3778 const HloInstruction* xla_while, in CheckWhileBuffersShareAllocation() argument
3781 xla_while->shape(), in CheckWhileBuffersShareAllocation()
3784 xla_while->while_condition()->parameter_instruction(0); in CheckWhileBuffersShareAllocation()
3785 const HloComputation* body = xla_while->while_body(); in CheckWhileBuffersShareAllocation()
3789 xla_while, xla_while->operand(0), index, buffer_assignment)); in CheckWhileBuffersShareAllocation()
[all …]
Dir_emitter_unnested.h195 Status HandleWhile(HloInstruction* xla_while) override;
/external/tensorflow/tensorflow/compiler/tf2xla/python/
Dxla.py478 while_loop = gen_xla_ops.xla_while