From 596a7f0da818eb7c0cd27130c2d7133e4a8ca4af Mon Sep 17 00:00:00 2001 From: Mark Heffernan Date: Mon, 9 Jul 2018 14:53:47 -0700 Subject: Add WithToken variants to Send/Recv/Infeed/Outfeed ops in XLA builder. In HLO side-effecting ops now take and produce a token-shaped value used for ordering. This CL exposes the tokens in the XLA builder interface by adding builder variants suffixed with "WithToken". These variants have an extra token-shaped operand and produce a token, sometimes as a element in the tuple-shaped output of the op. Also, at the HLO level, add new method CreateToken for creating an instruction which produces token-shaped value out of thin air using an AfterAll instruction. This is explicitly broken out from CreateAfterAll to facilitate remove of zero-operand AfterAll instructions (b/110532604). PiperOrigin-RevId: 203834271 --- tensorflow/compiler/xla/service/conditional_simplifier_test.cc | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) (limited to 'tensorflow/compiler/xla/service/conditional_simplifier_test.cc') diff --git a/tensorflow/compiler/xla/service/conditional_simplifier_test.cc b/tensorflow/compiler/xla/service/conditional_simplifier_test.cc index 834878426f..c43a31b167 100644 --- a/tensorflow/compiler/xla/service/conditional_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/conditional_simplifier_test.cc @@ -119,8 +119,7 @@ TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsSend) { ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional); auto* true_computation = conditional->true_computation(); - auto* token = - true_computation->AddInstruction(HloInstruction::CreateAfterAll({})); + auto* token = true_computation->AddInstruction(HloInstruction::CreateToken()); auto* send = true_computation->AddInstruction(HloInstruction::CreateSend( true_computation->AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))), @@ -135,8 +134,7 @@ TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsRecv) { ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional); auto* true_computation = conditional->true_computation(); - auto* token = - true_computation->AddInstruction(HloInstruction::CreateAfterAll({})); + auto* token = true_computation->AddInstruction(HloInstruction::CreateToken()); auto* recv = true_computation->AddInstruction(HloInstruction::CreateRecv( ShapeUtil::MakeShape(F32, {1}), token, /*channel_id=*/0)); true_computation->AddInstruction(HloInstruction::CreateRecvDone(recv)); @@ -148,8 +146,7 @@ TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsNonRemovableInstruction) { auto* conditional = computation->root_instruction(); ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional); auto* false_computation = conditional->false_computation(); - auto token = - false_computation->AddInstruction(HloInstruction::CreateAfterAll({})); + auto token = false_computation->AddInstruction(HloInstruction::CreateToken()); false_computation->AddInstruction(HloInstruction::CreateInfeed( ShapeUtil::MakeShape(F32, {1}), token, "config")); EXPECT_FALSE(ConditionalSimplifier().Run(&module()).ValueOrDie()); -- cgit v1.2.3