Home
last modified time | relevance | path

Searched refs:async_start (Results 1 – 16 of 16) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dasync_op_canonicalizer_test.cc64 HloInstruction* async_start = FindInstruction(module.get(), "async-start"); in TEST_F() local
68 EXPECT_EQ(async_start->async_group_id(), 0); in TEST_F()
71 EXPECT_EQ(async_start->async_wrapped_computation(), in TEST_F()
73 EXPECT_EQ(async_start->async_wrapped_computation(), in TEST_F()
127 HloInstruction* async_start = FindInstruction(module.get(), "async-start"); in TEST_F() local
131 EXPECT_EQ(async_start->async_group_id(), 0); in TEST_F()
134 EXPECT_EQ(async_start->async_wrapped_computation(), in TEST_F()
136 EXPECT_EQ(async_start->async_wrapped_computation(), in TEST_F()
Dasync_collective_creator.cc162 HloInstruction* async_start = async_done->mutable_operand(0); in Run() local
163 replaced_pairs[ata] = ReplacedAsync{async_start, async_done}; in Run()
Dtuple_points_to_analysis.cc321 Status TuplePointsToAnalysis::HandleAsyncStart(HloInstruction* async_start) { in HandleAsyncStart() argument
323 PointsToSet& points_to_set = CreateEmptyPointsToSet(async_start); in HandleAsyncStart()
329 GetPointsToSet(async_start->operand(target_index.at(1))); in HandleAsyncStart()
338 &logical_buffer_analysis_->GetBuffer(async_start, target_index)); in HandleAsyncStart()
Ddfs_hlo_visitor_with_default.h229 Status HandleAsyncStart(HloInstructionPtr async_start) override { in HandleAsyncStart() argument
230 return DefaultAction(async_start); in HandleAsyncStart()
Dhlo_instruction_test.cc795 auto* async_start = async_done->operand(0); in TEST_F() local
797 EXPECT_EQ(async_start->shape().tuple_shapes_size(), 3); in TEST_F()
798 EXPECT_EQ(async_start->async_execution_thread(), "parallel_thread"); in TEST_F()
800 EXPECT_TRUE(ShapeUtil::Equal(async_start->shape().tuple_shapes(2), in TEST_F()
802 EXPECT_EQ(async_start->async_wrapped_computation()->execution_thread(), in TEST_F()
806 EXPECT_THAT(async_start->operands(), ElementsAre(constant1, constant2)); in TEST_F()
807 EXPECT_THAT(constant1->users(), ElementsAre(async_start)); in TEST_F()
808 EXPECT_THAT(constant2->users(), ElementsAre(async_start)); in TEST_F()
1722 HloInstruction* async_start = in TEST_F() local
1729 s_tuple, async_start, async_computation.get(), in TEST_F()
Dhlo_dataflow_analysis.cc605 HloInstruction* async_start) { in UpdateAsyncStartValueSet() argument
606 CHECK_EQ(async_start->opcode(), HloOpcode::kAsyncStart); in UpdateAsyncStartValueSet()
609 for (int64_t i = 0; i < async_start->operand_count(); ++i) { in UpdateAsyncStartValueSet()
610 const HloInstruction* operand = async_start->operand(i); in UpdateAsyncStartValueSet()
621 HloValueSet& value_set = GetValueSet(async_start, output_index); in UpdateAsyncStartValueSet()
631 async_start->async_wrapped_computation()->root_instruction(); in UpdateAsyncStartValueSet()
642 HloValueSet& value_set = GetValueSet(async_start, output_index); in UpdateAsyncStartValueSet()
Dhlo_verifier.cc1454 Status ShapeVerifier::HandleAsyncStart(HloInstruction* async_start) { in HandleAsyncStart() argument
1456 CheckAsyncOpComputationShapes(async_start, async_start->shape())); in HandleAsyncStart()
1457 TF_RETURN_IF_ERROR(CheckAsyncOpComputationThreadName(async_start)); in HandleAsyncStart()
1458 const Shape& param_shape = async_start->shape().tuple_shapes(0); in HandleAsyncStart()
1459 for (int i = 0; i < async_start->operand_count(); ++i) { in HandleAsyncStart()
1460 if (param_shape.tuple_shapes(i) != async_start->operand(i)->shape()) { in HandleAsyncStart()
1464 HloOpcodeString(async_start->opcode()), i, in HandleAsyncStart()
1465 async_start->operand(i)->shape().ToString(/*print_layout=*/true), in HandleAsyncStart()
Dhlo_dataflow_analysis.h270 bool UpdateAsyncStartValueSet(HloInstruction* async_start);
Dhlo_computation.cc780 HloInstruction* async_start = AddInstruction(HloInstruction::CreateAsyncStart( in CreateAsyncInstructions() local
785 root->shape(), async_start, async_computation, in CreateAsyncInstructions()
787 async_start->set_metadata(instruction->metadata()); in CreateAsyncInstructions()
788 async_start->CopyBackendConfigFrom(instruction); in CreateAsyncInstructions()
Dtuple_points_to_analysis.h249 Status HandleAsyncStart(HloInstruction* async_start) override;
Dhlo_cost_analysis.h100 Status HandleAsyncStart(const HloInstruction* async_start) override;
Dhlo_dataflow_analysis_test.cc1090 const HloInstruction* async_start = in TEST_P() local
1097 async_start->async_wrapped_instruction(); in TEST_P()
1099 EXPECT_TRUE(analysis.ValueIsDefinedAt(async_start, /*index=*/{})); in TEST_P()
1100 EXPECT_FALSE(analysis.ValueIsDefinedAt(async_start, /*index=*/{0, 0})); in TEST_P()
1101 EXPECT_FALSE(analysis.ValueIsDefinedAt(async_start, /*index=*/{1})); in TEST_P()
1102 EXPECT_THAT(HloValuesAt(async_start, {1}), in TEST_P()
1105 EXPECT_TRUE(analysis.ValueIsDefinedAt(async_start, /*index=*/{2})); in TEST_P()
1106 EXPECT_THAT(HloValuesAt(async_start, /*index=*/{0, 0}), in TEST_P()
Dhlo_evaluator.h301 Status HandleAsyncStart(HloInstruction* async_start) override;
Dhlo_evaluator.cc3034 Status HloEvaluator::HandleAsyncStart(HloInstruction* async_start) { in HandleAsyncStart() argument
3036 arg_literals.reserve(async_start->operands().size()); in HandleAsyncStart()
3037 for (auto operand : async_start->operands()) { in HandleAsyncStart()
3047 embedded_evaluator.Evaluate(*async_start->async_wrapped_computation(), in HandleAsyncStart()
3050 evaluated_[async_start] = Literal(async_start->shape()); in HandleAsyncStart()
3053 TF_RETURN_IF_ERROR(evaluated_[async_start].CopyFrom( in HandleAsyncStart()
3058 TF_RETURN_IF_ERROR(evaluated_[async_start].MoveFrom( in HandleAsyncStart()
Dhlo_verifier.h200 Status HandleAsyncStart(HloInstruction* async_start) override;
Dhlo_cost_analysis.cc487 Status HloCostAnalysis::HandleAsyncStart(const HloInstruction* async_start) { in HandleAsyncStart() argument
490 ProcessSubcomputation(async_start->called_computations()[0])); in HandleAsyncStart()