From e730b261f9028b2f3430461b82c30c86b9ece22f Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 9 Oct 2018 06:58:06 -0700 Subject: Automated rollback of commit 375c109659d2d0e6265447dffdeb460693b3cccf PiperOrigin-RevId: 216350134 --- tensorflow/compiler/xla/service/BUILD | 21 --- .../compiler/xla/service/buffer_assignment.cc | 34 ++-- tensorflow/compiler/xla/service/buffer_value.h | 3 - tensorflow/compiler/xla/service/copy_insertion.cc | 85 +--------- .../compiler/xla/service/copy_insertion_test.cc | 183 -------------------- tensorflow/compiler/xla/service/hlo.proto | 29 ---- .../compiler/xla/service/hlo_alias_analysis.cc | 46 +----- .../xla/service/hlo_alias_analysis_test.cc | 175 -------------------- .../compiler/xla/service/hlo_dataflow_analysis.cc | 2 +- .../xla/service/hlo_input_output_alias_config.cc | 172 ------------------- .../xla/service/hlo_input_output_alias_config.h | 101 ----------- .../service/hlo_input_output_alias_config_test.cc | 184 --------------------- tensorflow/compiler/xla/service/hlo_module.cc | 9 - tensorflow/compiler/xla/service/hlo_module.h | 14 -- tensorflow/compiler/xla/service/hlo_verifier.cc | 2 - tensorflow/compiler/xla/shape_util.h | 2 +- 16 files changed, 25 insertions(+), 1037 deletions(-) delete mode 100644 tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc delete mode 100644 tensorflow/compiler/xla/service/hlo_input_output_alias_config.h delete mode 100644 tensorflow/compiler/xla/service/hlo_input_output_alias_config_test.cc diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index 26ebb88e96..2b292ed053 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -294,7 +294,6 @@ cc_library( srcs = [ "dfs_hlo_visitor.cc", "hlo_computation.cc", - "hlo_input_output_alias_config.cc", "hlo_instruction.cc", "hlo_instructions.cc", "hlo_module.cc", @@ -309,7 +308,6 @@ cc_library( "hlo_clone_context.h", "hlo_computation.h", "hlo_domain_metadata.h", - "hlo_input_output_alias_config.h", "hlo_instruction.h", "hlo_instructions.h", "hlo_module.h", @@ -1270,25 +1268,6 @@ tf_cc_test( ], ) -tf_cc_test( - name = "hlo_input_output_alias_config_test", - srcs = ["hlo_input_output_alias_config_test.cc"], - deps = [ - ":hlo", - ":hlo_dce", - ":hlo_memory_scheduler", - ":hlo_ordering", - ":hlo_parser", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:xla_internal_test_main", - "//tensorflow/core:test", - "@com_google_absl//absl/algorithm:container", - ], -) - cc_library( name = "hlo_memory_scheduler", srcs = ["hlo_memory_scheduler.cc"], diff --git a/tensorflow/compiler/xla/service/buffer_assignment.cc b/tensorflow/compiler/xla/service/buffer_assignment.cc index d5d6a044a8..2c2d1626c2 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment.cc @@ -239,7 +239,7 @@ BufferAllocation::Slice BufferAllocation::GetSlice( void BufferAllocation::AddAssignment(const LogicalBuffer& buffer, int64 offset, int64 size) { - VLOG(4) << "Trying to add " << buffer << " to allocation #" << index(); + VLOG(4) << "Trying to add " << buffer << " to " << this; CHECK(assigned_buffers_.count(&buffer) == 0) << "LogicalBuffer " << buffer << " already assigned to allocation " << index_; @@ -784,6 +784,21 @@ bool BufferAssigner::MaybeAssignBuffer(BufferAllocation* allocation, } } + if (allow_input_output_aliasing_ && allocation->maybe_live_out()) { + const HloComputation* entry_computation = + assignment->module_->entry_computation(); + for (auto param : entry_computation->parameter_instructions()) { + for (auto& param_buffer : + assignment->points_to_analysis().GetBuffersDefinedByInstruction( + param)) { + if (assignment->liveness().MayInterfere(*param_buffer, buffer)) { + VLOG(4) << "Can't assign: Parameter interference with result"; + return false; + } + } + } + } + // If the buffer is live out of the computation then it should only be // assigned a buffer which exactly fits the result to avoid wasting memory // (result buffers can have arbitrary lifetimes). @@ -1419,28 +1434,13 @@ BufferAssigner::MergeColocatedBufferSets( // Builds sets of buffers in 'colocated_buffer_sets' which should be colocated // in the same allocation (currently just supports kWhile, kCall, and -// kConditional and input output aliasing). +// kConditional). void BufferAssigner::BuildColocatedBufferSets( const HloModule* module, const BufferLiveness& buffer_liveness, const LogicalBuffer::SizeFunction& buffer_size, std::vector* colocated_buffer_sets) { const TuplePointsToAnalysis& points_to_analysis = buffer_liveness.points_to_analysis(); - - // Set up colocated buffer set for input and output. - module->input_output_alias_config().ForEachAlias( - [&](const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index) { - std::vector colocated_set; - AddBufferToColocatedSet(module->entry_computation()->root_instruction(), - output_index, points_to_analysis, - &colocated_set); - AddBufferToColocatedSet( - module->entry_computation()->parameter_instruction(param_number), - param_index, points_to_analysis, &colocated_set); - AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); - }); - for (const HloComputation* computation : module->MakeComputationPostOrder()) { if (computation->IsFusionComputation()) { continue; diff --git a/tensorflow/compiler/xla/service/buffer_value.h b/tensorflow/compiler/xla/service/buffer_value.h index 11d8abc5ba..69b3646356 100644 --- a/tensorflow/compiler/xla/service/buffer_value.h +++ b/tensorflow/compiler/xla/service/buffer_value.h @@ -141,9 +141,6 @@ class BufferValue { // operator< is required for std::set. bool operator<(const BufferValue& other) const { return id_ < other.id_; } - bool operator==(const BufferValue& other) const { return id_ == other.id_; } - bool operator!=(const BufferValue& other) const { return id_ != other.id_; } - virtual string ToString() const = 0; // TODO(lauj) rename LogicalBufferProto to BufferValueProto. diff --git a/tensorflow/compiler/xla/service/copy_insertion.cc b/tensorflow/compiler/xla/service/copy_insertion.cc index cfe025fdd1..f35324aa35 100644 --- a/tensorflow/compiler/xla/service/copy_insertion.cc +++ b/tensorflow/compiler/xla/service/copy_insertion.cc @@ -40,12 +40,10 @@ namespace { using absl::StrAppend; -bool IsReadonlyEntryParameterValue(const HloValue& value) { +bool IsEntryParameterValue(const HloValue& value) { const HloComputation* computation = value.defining_instruction()->parent(); return value.defining_instruction()->opcode() == HloOpcode::kParameter && - computation == computation->parent()->entry_computation() && - !computation->parent()->input_output_alias_config().ParameterHasAlias( - value.defining_instruction()->parameter_number()); + computation == computation->parent()->entry_computation(); } bool IsConstantValue(const HloValue& value) { @@ -53,7 +51,7 @@ bool IsConstantValue(const HloValue& value) { } bool ValueIsReadOnly(const HloValue& value) { - return IsConstantValue(value) || IsReadonlyEntryParameterValue(value); + return IsConstantValue(value) || IsEntryParameterValue(value); } // Data structure describing the action which should be taken on parts of a @@ -334,81 +332,6 @@ Status AddCopiesForConditional(const HloAliasAnalysis& alias_analysis, return Status::OK(); } -// Conservatively adds copies before root instruction of entry computation and -// each aliased parameter to resolve interference of aliased input and output -// buffer. We later rely on the CopyRemover to drop the unnecessary ones. -Status AddCopiesForAliasedInputOutputs(HloModule* module) { - HloComputation* entry = module->entry_computation(); - HloInstruction* root = entry->root_instruction(); - - ShapeTree output_indices_to_copy(root->shape()); - std::vector> copied_parameters; - bool has_alias = false; - for (auto* param : entry->parameter_instructions()) { - bool param_has_alias = false; - ShapeTree param_indices_to_copy(param->shape()); - - module->input_output_alias_config().ForEachAlias( - [&](const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index) { - if (param_number == param->parameter_number()) { - param_has_alias = true; - *(param_indices_to_copy.mutable_element(param_index)) = true; - *(output_indices_to_copy.mutable_element(output_index)) = true; - } - }); - - if (!param_has_alias) { - continue; - } - - has_alias = true; - // Store a snapshot of users before DeepCopyInstruction, as - // DeepCopyInstruction introduces new users of the instruction. - std::vector users = param->users(); - ShapeTree param_copy_tree(param->shape(), - /*init_value=*/nullptr); - TF_ASSIGN_OR_RETURN(HloInstruction * copied, - entry->DeepCopyInstruction( - param, ¶m_indices_to_copy, ¶m_copy_tree)); - for (HloInstruction* user : users) { - TF_RETURN_IF_ERROR(param->ReplaceUseWith(user, copied)); - } - - copied_parameters.push_back(param_copy_tree); - } - - if (!has_alias) { - return Status::OK(); - } - - // Add copies before root instruction. - ShapeTree output_copy_tree(root->shape(), - /*init_value=*/nullptr); - - TF_ASSIGN_OR_RETURN(HloInstruction * root_copied, - root->parent()->DeepCopyInstruction( - root, &output_indices_to_copy, &output_copy_tree)); - - // Add control dependencies between the input/output copies. - TF_RETURN_IF_ERROR(module->input_output_alias_config().ForEachAliasWithStatus( - [&](const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& input_index) -> Status { - HloInstruction* from = - copied_parameters[param_number].element(input_index); - HloInstruction* to = output_copy_tree.element(output_index); - - TF_RET_CHECK(from != nullptr); - TF_RET_CHECK(to != nullptr); - TF_RETURN_IF_ERROR(from->AddControlDependencyTo(to)); - return Status::OK(); - })); - - entry->set_root_instruction(root_copied); - - return Status::OK(); -} - // Removes any control dependencies to or from the given instruction. Status StripControlDependenciesFrom(HloInstruction* instruction) { while (!instruction->control_successors().empty()) { @@ -1030,8 +953,6 @@ Status CopyInsertion::AddCopiesToResolveInterference(HloModule* module) { } } } - - TF_RETURN_IF_ERROR(AddCopiesForAliasedInputOutputs(module)); return Status::OK(); } diff --git a/tensorflow/compiler/xla/service/copy_insertion_test.cc b/tensorflow/compiler/xla/service/copy_insertion_test.cc index 3096206c34..892d0d7b54 100644 --- a/tensorflow/compiler/xla/service/copy_insertion_test.cc +++ b/tensorflow/compiler/xla/service/copy_insertion_test.cc @@ -1351,189 +1351,6 @@ TEST_F(CopyInsertionTest, SwizzlingWhile) { EXPECT_THAT(xla_while->operand(0), op::Tuple(op::Copy(), op::Copy())); } -TEST_F(CopyInsertionTest, CrossingParameters) { - // Test a case where two parameters' dataflow cross with each other while - // input and output are aliased with same index: - // - // (p0 , p1) - // | \ /| - // | \ / | - // alias X alias - // | / \ | - // | / \| - // (p1 , p0) - auto module = CreateNewModule(); - const Shape tuple_shape = - ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); - - auto builder = HloComputation::Builder(TestName()); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "0")); - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 1)); - builder.AddInstruction(HloInstruction::CreateTuple({gte1, gte0})); - module->AddEntryComputation(builder.Build()); - ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); - ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1})); - InsertCopies(module.get()); - - EXPECT_EQ(CountCopies(*module), 4); -} - -TEST_F(CopyInsertionTest, ParametersAliasing) { - // Test a case where two parameters' dataflow don't interfere with each other - // while aliased. - // - // (p0 , p1) - // | | - // | | - // alias alias - // | | - // | | - // (p0 , p1) - auto module = CreateNewModule(); - const Shape tuple_shape = - ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); - - auto builder = HloComputation::Builder(TestName()); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 1)); - builder.AddInstruction(HloInstruction::CreateTuple({gte0, gte1})); - module->AddEntryComputation(builder.Build()); - ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); - ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1})); - InsertCopies(module.get()); - - EXPECT_THAT(module->entry_computation()->root_instruction(), - op::Tuple(op::Copy(op::GetTupleElement(param, 0)), - op::Copy(op::GetTupleElement(param, 1)))); - - EXPECT_EQ(CountCopies(*module), 2); -} - -TEST_F(CopyInsertionTest, ParameterWithPartialAliasing) { - // Test a case where one parameter is aliased with result while another one - // isn't. - // - // (p0 , p1) - // | | - // | | - // alias | - // | | - // | | - // (p0 , p1) - auto module = CreateNewModule(); - const Shape tuple_shape = - ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); - - auto builder = HloComputation::Builder(TestName()); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 1)); - builder.AddInstruction(HloInstruction::CreateTuple({gte0, gte1})); - module->AddEntryComputation(builder.Build()); - ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); - InsertCopies(module.get()); - - EXPECT_THAT(module->entry_computation()->root_instruction(), - op::Tuple(op::Copy(op::GetTupleElement(param, 0)), - op::Copy(op::GetTupleElement(param, 1)))); - - EXPECT_EQ(CountCopies(*module), 2); -} - -TEST_F(CopyInsertionTest, ParameterAndParallelOpsWithPartialAliasing) { - // Test a case where one parameter is aliased with result while another one - // isn't. - // - // +-- (p0 , p1) - // | | | - // | | | - // alias Negate Negate - // | | | - // | | | - // +-- (p0 , p1) - auto module = CreateNewModule(); - const Shape tuple_shape = - ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); - - auto builder = HloComputation::Builder(TestName()); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 1)); - - auto negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(scalar_shape_, HloOpcode::kNegate, gte0)); - - auto negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(scalar_shape_, HloOpcode::kNegate, gte1)); - builder.AddInstruction(HloInstruction::CreateTuple({negate0, negate1})); - module->AddEntryComputation(builder.Build()); - ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); - InsertCopies(module.get()); - - EXPECT_EQ(CountCopies(*module), 0); -} - -TEST_F(CopyInsertionTest, ParameterAndOpsWithPartialAliasing) { - // Test a case where one parameter is aliased with result while another one - // isn't. - // - // +-- (p0 , p1) - // | | | - // | | | - // alias Negate Negate - // | | | - // | Add----+ - // | | | - // +-- (p0 , p1) - auto module = CreateNewModule(); - const Shape tuple_shape = - ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); - - auto builder = HloComputation::Builder(TestName()); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 1)); - - auto negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(scalar_shape_, HloOpcode::kNegate, gte0)); - - auto negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(scalar_shape_, HloOpcode::kNegate, gte1)); - - auto add = builder.AddInstruction(HloInstruction::CreateBinary( - scalar_shape_, HloOpcode::kAdd, negate0, negate1)); - builder.AddInstruction(HloInstruction::CreateTuple({add, negate1})); - module->AddEntryComputation(builder.Build()); - ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); - InsertCopies(module.get()); - - EXPECT_EQ(CountCopies(*module), 0); -} - TEST_F(CopyInsertionTest, SwizzlingWhileWithOneOp) { // Test a while instruction with a body which permutes its tuple parameter // elements and applies one operation to one of the elements. The addition of diff --git a/tensorflow/compiler/xla/service/hlo.proto b/tensorflow/compiler/xla/service/hlo.proto index 82c8fb1904..a0eb9e6ddc 100644 --- a/tensorflow/compiler/xla/service/hlo.proto +++ b/tensorflow/compiler/xla/service/hlo.proto @@ -225,32 +225,6 @@ message HloScheduleProto { map sequences = 1; } -message HloInputOutputAliasProto { - // The following proto describes a pair of aliased an input - // (described by parameter number and a ShapeIndex of the parameter) - // and an output (described by a ShapeIndex of the root - // instruction). For example: - // - // entry = { - // output_shape_index={1}, - // parameter_number=0, - // parameter_shape_index={1, 2}, - // } - // - // This entry indicates that the first paremter's {1, 2} element is - // aliased with the {1} element of the root instruction. - message AliasEntryProto { - // ShapeIndex of the root hlo. - repeated int64 output_shape_index = 1; - // Number of the parameter in entry computation. - int64 parameter_number = 2; - // ShapeIndex of the parameter instruction. - repeated int64 parameter_shape_index = 3; - } - - repeated AliasEntryProto entries = 1; -} - // Serialization of HloModule. message HloModuleProto { string name = 1; @@ -269,9 +243,6 @@ message HloModuleProto { // The schedule for this module. HloScheduleProto schedule = 7; - - // Describes alias information between inputs and outputs. - HloInputOutputAliasProto input_output_alias = 8; } // Serialization of LogicalBuffer. diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc index cf8e6594cb..c3da12e273 100644 --- a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc @@ -59,9 +59,8 @@ class BufferValueMap { // construction process. using BufferNumber = int64; - explicit BufferValueMap(HloModule* module, - const HloDataflowAnalysis& dataflow) - : module_(module), dataflow_(dataflow) { + explicit BufferValueMap(const HloDataflowAnalysis& dataflow) + : dataflow_(dataflow) { buffers_.reserve(dataflow_.values().size()); value_to_buffer_number_.reserve(dataflow_.values().size()); for (const HloValue* value : dataflow_.values()) { @@ -172,42 +171,6 @@ class BufferValueMap { return value_to_buffer_number_.at(&value); } - void ComputeInputOutputAliasedBuffers( - const HloValue& value, std::vector* aliased_buffers) { - // Get parameter value from an aliased_input object. - const auto get_parameter_value = - [this](const std::pair& aliased_input) - -> const HloValue& { - int64 param_number = aliased_input.first; - const ShapeIndex& param_index = aliased_input.second; - return dataflow_.GetUniqueValueAt( - module_->entry_computation()->parameter_instruction(param_number), - param_index); - }; - - // If the value shows up in a root instruction, alias it with parameter - // intruction. - for (const HloPosition& pos : value.positions()) { - if (pos.instruction == module_->entry_computation()->root_instruction()) { - ShapeIndex output_index = pos.index; - - auto aliased_input = - module_->input_output_alias_config().GetAliasedParameter( - output_index); - if (aliased_input) { - aliased_buffers->push_back( - GetBufferForValue(get_parameter_value(*aliased_input))); - } - } - } - - // If the value is parameter instruction itself, alias it with itself. - if (value.instruction()->opcode() == HloOpcode::kParameter && - value.instruction()->parent() == module_->entry_computation()) { - aliased_buffers->push_back(GetBufferForValue(value)); - } - } - void ComputeWhileAliasedBuffers(const HloValue& value, std::vector* aliased_buffers) { VLOG(3) << "Compute kWhile aliases"; @@ -315,7 +278,6 @@ class BufferValueMap { VLOG(2) << "Use of value " << value.ToShortString() << ": " << use; } std::vector aliased_buffers; - ComputeInputOutputAliasedBuffers(value, &aliased_buffers); ComputeWhileAliasedBuffers(value, &aliased_buffers); ComputeConditionalAliasedBuffers(value, &aliased_buffers); // Uniquify aliased buffers. @@ -326,8 +288,6 @@ class BufferValueMap { return aliased_buffers; } - HloModule* module_; - // Dataflow analysis used to construct the buffer map. const HloDataflowAnalysis& dataflow_; @@ -501,7 +461,7 @@ StatusOr> HloAliasAnalysis::Run( /*bitcast_defines_value=*/false, fusion_can_share_buffer)); - BufferValueMap buffer_map(module, alias_analysis->dataflow_analysis()); + BufferValueMap buffer_map(alias_analysis->dataflow_analysis()); buffer_map.MergeAliasedBuffers(); // Create a vector of HloBuffers, one for each set of values in the diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc index 5c8d97b2d1..0cd0ab36fc 100644 --- a/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc @@ -217,181 +217,6 @@ TEST_F(HloAliasAnalysisTest, NondistinctTuple) { EXPECT_FALSE(AnyValuesInSameBufferInterfere()); } -TEST_F(HloAliasAnalysisTest, ParametersWithAliasing) { - const Shape tuple_shape = - ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); - - auto builder = HloComputation::Builder(TestName()); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 1)); - - auto negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(scalar_shape_, HloOpcode::kNegate, gte0)); - auto negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(scalar_shape_, HloOpcode::kNegate, gte1)); - - auto tuple = - builder.AddInstruction(HloInstruction::CreateTuple({negate0, negate1})); - module_->AddEntryComputation(builder.Build()); - TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); - TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1})); - - // Cannot alias an output twice. - ASSERT_IS_NOT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{0})); - - const HloAliasAnalysis& analysis = RunAnalysis(); - - EXPECT_EQ(analysis.GetUniqueBufferAt(gte0), - analysis.GetUniqueBufferAt(tuple, /*index=*/{0})); - - EXPECT_EQ(analysis.GetUniqueBufferAt(gte1), - analysis.GetUniqueBufferAt(tuple, /*index=*/{1})); -} - -TEST_F(HloAliasAnalysisTest, ParametersWithCrossAliasing) { - // parameter 0 aliased with output 1 and parameter 1 aliased with output 0. - // - // (p0 , p1) - // \ / - // \ / - // alias X - // / \ - // / \ - // (p0 , p1) - const Shape tuple_shape = - ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); - - auto builder = HloComputation::Builder(TestName()); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, param, 1)); - auto tuple = - builder.AddInstruction(HloInstruction::CreateTuple({gte0, gte1})); - module_->AddEntryComputation(builder.Build()); - TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{1})); - TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{0})); - - // Cannot alias an output twice. - ASSERT_IS_NOT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1})); - - const HloAliasAnalysis& analysis = RunAnalysis(); - - // Every Ops in this graph are aliased with each other. - EXPECT_EQ(analysis.GetUniqueBufferAt(gte0), - analysis.GetUniqueBufferAt(tuple, /*index=*/{0})); - EXPECT_EQ(analysis.GetUniqueBufferAt(gte0), - analysis.GetUniqueBufferAt(tuple, /*index=*/{1})); - - EXPECT_EQ(analysis.GetUniqueBufferAt(gte1), - analysis.GetUniqueBufferAt(tuple, /*index=*/{0})); - EXPECT_EQ(analysis.GetUniqueBufferAt(gte1), - analysis.GetUniqueBufferAt(tuple, /*index=*/{1})); -} - -TEST_F(HloAliasAnalysisTest, InputOutputAliasingWithWhile) { - // Test a simple single while instruction can be aliased with input and output - // of the computation. - // - // body((F32[], F32[]) %tuple_param): - // %add = Add(%tuple_param{0}, %tuple_param{1}) - // return Tuple(%tuple_param{0}, %add) - // - // condition((F32[], F32[]) %tuple_param): - // return Constant(false) - // - // entry: - // %param1 = param1 - // %while = While(%param1, body, condition) - // %while_1 = GTE(%while, 0) - // %while_2 = GTE(%while, 1) - // %negate_1 = Negate(%while_1) - // %negate_2 = Negate(%while_2) - // return Tuple(negate_1, negate_2) - // - const Shape tuple_shape = - ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); - - // Element 0 passes transparently through the body. - auto body_builder = HloComputation::Builder("body"); - auto body_param = body_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "param")); - auto body_element_0 = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 0)); - auto body_element_1 = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 1)); - auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( - scalar_shape_, HloOpcode::kAdd, body_element_0, body_element_1)); - auto body_tuple = body_builder.AddInstruction( - HloInstruction::CreateTuple({body_element_0, add})); - HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); - - // Condition computation trivially returns a constant "false". - auto cond_builder = HloComputation::Builder("condition"); - auto cond_param = cond_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "param")); - cond_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); - HloComputation* condition = - module_->AddEmbeddedComputation(cond_builder.Build()); - - auto builder = HloComputation::Builder(TestName()); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - - auto xla_while = builder.AddInstruction( - HloInstruction::CreateWhile(tuple_shape, condition, body, param)); - auto while_element_1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, xla_while, 0)); - auto while_element_2 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape_, xla_while, 1)); - auto negate_1 = builder.AddInstruction(HloInstruction::CreateUnary( - scalar_shape_, HloOpcode::kNegate, while_element_1)); - auto negate_2 = builder.AddInstruction(HloInstruction::CreateUnary( - scalar_shape_, HloOpcode::kNegate, while_element_2)); - auto tuple = - builder.AddInstruction(HloInstruction::CreateTuple({negate_1, negate_2})); - module_->AddEntryComputation(builder.Build()); - TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); - TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1})); - - const HloAliasAnalysis& analysis = RunAnalysis(); - - EXPECT_THAT( - GetValuesInBuffer(analysis.GetUniqueBufferAt(xla_while, /*index=*/{1})), - UnorderedElementsAre(GetValueDefinedAt(param, {1}), - GetValueDefinedAt(xla_while, /*index=*/{1}), - GetValueDefinedAt(body_param, {1}), - GetValueDefinedAt(cond_param, {1}), - GetValueDefinedAt(add), - GetValueDefinedAt(negate_2))); - - EXPECT_THAT( - analysis.GetUniqueBufferAt(xla_while, /*index=*/{1}).ComputePositions(), - UnorderedElementsAre( - HloPosition{param, {1}}, HloPosition{xla_while, {1}}, - HloPosition{while_element_2, {}}, HloPosition{body_param, {1}}, - HloPosition{body_element_1, {}}, HloPosition{add, {}}, - HloPosition{body_tuple, {1}}, HloPosition{tuple, {1}}, - HloPosition{cond_param, {1}}, HloPosition{negate_2, {}})); - - EXPECT_FALSE(AnyValuesInSameBufferInterfere()); -} - TEST_F(HloAliasAnalysisTest, SingleCall) { // Test a single call of a subcomputation. The subcomputation adds its two // array-shaped parameters. diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc index f401eac016..c22adcdd8d 100644 --- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc @@ -126,7 +126,7 @@ bool HloDataflowAnalysis::ValueIsDefinedAt(const HloInstruction* instruction, const HloValue& HloDataflowAnalysis::GetValueDefinedAt( const HloInstruction* instruction, const ShapeIndex& index) const { - CHECK(ValueIsDefinedAt(instruction, index)) << instruction->ToString(); + CHECK(ValueIsDefinedAt(instruction, index)); return GetUniqueValueAt(instruction, index); } diff --git a/tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc b/tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc deleted file mode 100644 index 9ad98e5038..0000000000 --- a/tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc +++ /dev/null @@ -1,172 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/xla/service/hlo_input_output_alias_config.h" -#include "tensorflow/compiler/xla/service/hlo_module.h" - -namespace xla { -Status HloInputOutputAliasConfig::SetUpAlias(const ShapeIndex& output_index, - int64 param_number, - const ShapeIndex& param_index) { - // Output can't be aliased with multiple parameters. - TF_RET_CHECK(!alias_.element(output_index)); - (*alias_.mutable_element(output_index)) = - std::make_pair(param_number, param_index); - return Status::OK(); -} - -HloInputOutputAliasProto HloInputOutputAliasConfig::ToProto() const { - HloInputOutputAliasProto result; - alias_.ForEachElement( - [&](const ShapeIndex& index, - const absl::optional>& data) { - if (data) { - HloInputOutputAliasProto::AliasEntryProto entry; - for (int64 i : index) { - entry.add_output_shape_index(i); - } - entry.set_parameter_number(data->first); - for (int64 i : data->second) { - entry.add_parameter_shape_index(i); - } - result.add_entries()->Swap(&entry); - } - }); - return result; -} - -StatusOr HloInputOutputAliasConfig::CreateFromProto( - const HloModule* module, const HloInputOutputAliasProto& proto) { - HloInputOutputAliasConfig result( - module->entry_computation()->root_instruction()->shape()); - for (const HloInputOutputAliasProto::AliasEntryProto& entry : - proto.entries()) { - ShapeIndex output_index(entry.output_shape_index().begin(), - entry.output_shape_index().end()); - - int64 param_number = entry.parameter_number(); - ShapeIndex param_index(entry.parameter_shape_index().begin(), - entry.parameter_shape_index().end()); - TF_RETURN_IF_ERROR( - result.SetUpAlias(output_index, param_number, param_index)); - } - - return result; -} - -string HloInputOutputAliasConfig::ToString() const { - std::vector pieces; - pieces.push_back("HloInputOutputAliasConfig"); - - ForEachAlias([&](const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index) { - pieces.push_back(absl::StrFormat( - " OutputIndex %s is aliased with parameter %lld at %s:", - output_index.ToString(), param_number, param_index.ToString())); - }); - - return absl::StrJoin(pieces, "\n"); -} - -bool HloInputOutputAliasConfig::ParameterHasAlias(int64 param_number) const { - bool output = false; - alias_.ForEachElement( - [&](const xla::ShapeIndex&, - absl::optional> alias) { - if (alias && alias->first == param_number) { - output = true; - } - }); - return output; -} - -absl::optional HloInputOutputAliasConfig::GetAliasedOutput( - int64 param_number, const ShapeIndex& param_index) const { - absl::optional output; - alias_.ForEachElement( - [&](const xla::ShapeIndex& output_index, - absl::optional> alias) { - if (alias && alias->first == param_number && - alias->second == param_index) { - output = output_index; - } - }); - return output; -} - -absl::optional> -HloInputOutputAliasConfig::GetAliasedParameter( - const ShapeIndex& output_index) const { - CHECK(ShapeUtil::IndexIsValid(alias_.shape(), output_index)); - return alias_.element(output_index); -} - -void HloInputOutputAliasConfig::ForEachAlias(AliasFn fn) const { - alias_.ForEachElement( - [&](const ShapeIndex& output_index, - absl::optional> aliased) { - if (aliased) { - fn(output_index, aliased->first, aliased->second); - } - }); -} - -Status HloInputOutputAliasConfig::ForEachAliasWithStatus( - AliasFnWithStatus fn) const { - return alias_.ForEachElementWithStatus( - [&](const ShapeIndex& output_index, - absl::optional> aliased) { - if (aliased) { - TF_RETURN_IF_ERROR(fn(output_index, aliased->first, aliased->second)); - } - return Status::OK(); - }); -} - -Status HloInputOutputAliasConfig::Verify(const HloModule& module) const { - std::vector> param_has_seen; - const HloComputation* entry = module.entry_computation(); - for (int64 i = 0; i < entry->num_parameters(); ++i) { - HloInstruction* param = entry->parameter_instruction(i); - param_has_seen.emplace_back(param->shape()); - } - return ForEachAliasWithStatus([&](const ShapeIndex& output_index, - int64 param_number, - const ShapeIndex& param_index) -> Status { - const HloInstruction* root = entry->root_instruction(); - - const Shape& param_shape = - entry->parameter_instruction(param_number)->shape(); - const Shape& output_shape = root->shape(); - TF_RET_CHECK(entry->num_parameters() > param_number); - TF_RET_CHECK(ShapeUtil::IndexIsValid(param_shape, param_index)); - TF_RET_CHECK(ShapeUtil::IndexIsValid(output_shape, output_index)); - - // Check each param_number and param_index pair only show up once. No - // input can be aliased with output buffers. - TF_RET_CHECK(param_has_seen[param_number].element(param_index) == false); - - *(param_has_seen[param_number].mutable_element(param_index)) = true; - - return Status::OK(); - }); -} - -std::ostream& operator<<(std::ostream& out, - const HloInputOutputAliasConfig& config) { - out << config.ToString(); - return out; -} -} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_input_output_alias_config.h b/tensorflow/compiler/xla/service/hlo_input_output_alias_config.h deleted file mode 100644 index 02c46f65c8..0000000000 --- a/tensorflow/compiler/xla/service/hlo_input_output_alias_config.h +++ /dev/null @@ -1,101 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_INPUT_OUTPUT_ALIAS_CONFIG_H_ -#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_INPUT_OUTPUT_ALIAS_CONFIG_H_ - -#include - -#include "absl/types/optional.h" -#include "tensorflow/compiler/xla/service/hlo.pb.h" -#include "tensorflow/compiler/xla/shape_tree.h" -#include "tensorflow/compiler/xla/shape_util.h" - -namespace xla { - -class HloModule; - -// This class specifies the alias map from output index to parameter number and -// parameter index in the entry computation. -class HloInputOutputAliasConfig { - public: - HloInputOutputAliasConfig() = default; - - explicit HloInputOutputAliasConfig(Shape shape) : alias_(shape) {} - - virtual ~HloInputOutputAliasConfig() = default; - - // Sets up alias config from `output_index` to `param_index` at - // `param_number`. - Status SetUpAlias(const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index); - - // Returns true if the given parameter is aliased with one of the output - // buffers. - bool ParameterHasAlias(int64 param_number) const; - - // (De)Serializes an HloInputOutoutAliasConfig to/from an - // HloInputOutoutAliasProto. - HloInputOutputAliasProto ToProto() const; - - static StatusOr CreateFromProto( - const HloModule* module, const HloInputOutputAliasProto& proto); - - // Returns the output index that the given parameter and parameter index is - // aliased with. A nullopt is returned if there is no output that is aliased - // with the parameter number and index. - absl::optional GetAliasedOutput( - int64 param_number, const ShapeIndex& param_index) const; - - // Returns the number of parameter and index of the parameter buffer that the - // given output buffer index is aliased with. A nullopt is returned if there - // is no parameter is aliased with the specific output. - absl::optional> GetAliasedParameter( - const ShapeIndex& output_index) const; - - using AliasFn = - std::function; - - // Iterates through each aliased output and input. - void ForEachAlias(AliasFn fn) const; - - using AliasFnWithStatus = - std::function; - - // Verifies that the given config is valid for the given module. - // Specifically, the config's input and output should be in-bound and size of - // the aliased buffers should match. - Status Verify(const HloModule& module) const; - - Status ForEachAliasWithStatus(AliasFnWithStatus fn) const; - - string ToString() const; - - private: - // A ShapeTree which indicates the list of buffers that's expected to be - // aliased. The key on this shape tree represents the output index. The value - // is a pair of parameter number and index into the buffer. If the value is - // nullopt, it means there is no parameter aliasing for this output. - ShapeTree>> alias_; -}; - -std::ostream& operator<<(std::ostream& out, - const HloInputOutputAliasConfig& config); - -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_INPUT_OUTPUT_ALIAS_CONFIG_H_ diff --git a/tensorflow/compiler/xla/service/hlo_input_output_alias_config_test.cc b/tensorflow/compiler/xla/service/hlo_input_output_alias_config_test.cc deleted file mode 100644 index 3b61ff04e6..0000000000 --- a/tensorflow/compiler/xla/service/hlo_input_output_alias_config_test.cc +++ /dev/null @@ -1,184 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/xla/service/hlo_input_output_alias_config.h" - -#include -#include - -#include "absl/algorithm/container.h" -#include "tensorflow/compiler/xla/service/hlo_computation.h" -#include "tensorflow/compiler/xla/service/hlo_dce.h" -#include "tensorflow/compiler/xla/service/hlo_instruction.h" -#include "tensorflow/compiler/xla/service/hlo_memory_scheduler.h" -#include "tensorflow/compiler/xla/service/hlo_opcode.h" -#include "tensorflow/compiler/xla/service/hlo_ordering.h" -#include "tensorflow/compiler/xla/service/hlo_parser.h" -#include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/tests/hlo_test_base.h" -#include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/lib/core/status_test_util.h" - -namespace xla { -namespace { -class HloInputOutputAliasConfigTest : public HloTestBase { - protected: - void expect_aliased(const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index, - const HloInputOutputAliasConfig& config) { - absl::optional aliased_output = - config.GetAliasedOutput(param_number, param_index); - - EXPECT_TRUE(aliased_output); - EXPECT_EQ(aliased_output.value(), output_index); - - absl::optional> aliased_param = - config.GetAliasedParameter(output_index); - - EXPECT_TRUE(aliased_param); - EXPECT_EQ(aliased_param.value(), std::make_pair(param_number, param_index)); - } - - void expect_not_aliased(const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index, - const HloInputOutputAliasConfig& config) { - absl::optional aliased_output = - config.GetAliasedOutput(param_number, param_index); - - EXPECT_FALSE(aliased_output && aliased_output == output_index); - - absl::optional> aliased_param = - config.GetAliasedParameter(output_index); - - EXPECT_FALSE(aliased_param && aliased_param->first == param_number && - aliased_param->second == param_index); - } -}; - -TEST_F(HloInputOutputAliasConfigTest, SimpleAliasing) { - const string module_str = R"( -HloModule TEST - -ENTRY main { - a = f32[] parameter(0) - b = f32[] parameter(1) - ROOT root = (f32[], f32[]) tuple(%a, %b) -} -)"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseHloString(module_str)); - - HloInputOutputAliasConfig config( - module->entry_computation()->root_instruction()->shape()); - - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{0}, /*param_number=*/1, - /*param_index=*/{})); - - expect_aliased(/*output_index=*/{0}, /*param_number=*/1, - /*param_index=*/{}, config); - - expect_not_aliased(/*output_index=*/{1}, /*param_number=*/1, - /*param_index=*/{}, config); - - expect_not_aliased(/*output_index=*/{0}, /*param_number=*/0, - /*param_index=*/{}, config); -} - -TEST_F(HloInputOutputAliasConfigTest, SimpleAliasingWithTupleInput) { - const string module_str = R"( -HloModule TEST - -ENTRY main { - param = (f32[], f32[]) parameter(0) - gte1 = f32[] get-tuple-element(%param), index=0 - gte2 = f32[] get-tuple-element(%param), index=1 - ROOT root = (f32[], f32[]) tuple(%gte1, %gte2) -} -)"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseHloString(module_str)); - - HloInputOutputAliasConfig config( - module->entry_computation()->root_instruction()->shape()); - - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{0}, /*param_number=*/0, - /*param_index=*/{0})); - - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{1}, /*param_number=*/0, - /*param_index=*/{1})); - - expect_aliased(/*output_index=*/{0}, /*param_number=*/0, - /*param_index=*/{0}, config); - - expect_aliased(/*output_index=*/{1}, /*param_number=*/0, - /*param_index=*/{1}, config); - - expect_not_aliased(/*output_index=*/{1}, /*param_number=*/1, - /*param_index=*/{}, config); - - expect_not_aliased(/*output_index=*/{0}, /*param_number=*/0, - /*param_index=*/{}, config); -} - -TEST_F(HloInputOutputAliasConfigTest, InputDoNotAliasTwice) { - const string module_str = R"( -HloModule TEST - -ENTRY main { - a = f32[] parameter(0) - b = f32[] parameter(1) - ROOT root = (f32[], f32[]) tuple(%a, %b) -} -)"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseHloString(module_str)); - - HloInputOutputAliasConfig config( - module->entry_computation()->root_instruction()->shape()); - - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{0}, /*param_number=*/0, - /*param_index=*/{})); - - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{1}, /*param_number=*/0, - /*param_index=*/{})); - - ASSERT_IS_NOT_OK(config.Verify(*module)); -} - -TEST_F(HloInputOutputAliasConfigTest, OutputDoNotAliasTwice) { - const string module_str = R"( -HloModule TEST - -ENTRY main { - a = f32[] parameter(0) - b = f32[] parameter(1) - ROOT root = (f32[], f32[]) tuple(%a, %b) -} -)"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseHloString(module_str)); - - HloInputOutputAliasConfig config( - module->entry_computation()->root_instruction()->shape()); - - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{0}, /*param_number=*/0, - /*param_index=*/{})); - - ASSERT_IS_NOT_OK(config.SetUpAlias(/*output_index=*/{0}, /*param_number=*/1, - /*param_index=*/{})); -} -} // namespace -} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc index 547f74a0ed..93e04eb3db 100644 --- a/tensorflow/compiler/xla/service/hlo_module.cc +++ b/tensorflow/compiler/xla/service/hlo_module.cc @@ -73,8 +73,6 @@ HloComputation* HloModule::AddComputationInternal( config_.SetDefaultComputationLayout( entry_computation_->ComputeProgramShape()); } - input_output_alias_config_ = HloInputOutputAliasConfig( - entry_computation_->root_instruction()->shape()); } if (uniquify_identifiers) { @@ -254,9 +252,6 @@ HloModuleProto HloModule::ToProto() const { if (has_schedule()) { *proto.mutable_schedule() = schedule().ToProto().ValueOrDie(); } - - *proto.mutable_input_output_alias() = input_output_alias_config().ToProto(); - return proto; } @@ -333,10 +328,6 @@ StatusOr> HloModule::CreateFromProto( } TF_RET_CHECK(module->entry_computation_ != nullptr); - TF_ASSIGN_OR_RETURN(module->input_output_alias_config_, - HloInputOutputAliasConfig::CreateFromProto( - module.get(), proto.input_output_alias())); - // Because we didn't uniquify the names or the ids, double-check that the // instruction and computation names and ids are unique from the proto. absl::flat_hash_set computation_names; diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h index 9b9dc3ba9f..735804e827 100644 --- a/tensorflow/compiler/xla/service/hlo_module.h +++ b/tensorflow/compiler/xla/service/hlo_module.h @@ -31,7 +31,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_clone_context.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" -#include "tensorflow/compiler/xla/service/hlo_input_output_alias_config.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/hlo_schedule.h" @@ -213,15 +212,6 @@ class HloModule { return result; } - // input_output_alias_config indicates the list of aliased buffers that are - // expected from the module. - HloInputOutputAliasConfig& input_output_alias_config() { - return input_output_alias_config_; - } - const HloInputOutputAliasConfig& input_output_alias_config() const { - return input_output_alias_config_; - } - // Returns the number of unique intruction ids given out. All ids up to // this point are guaranteed to be in the range [0..NumUniqueInstructionIds()) int NumUniqueInstructionIds() const { return next_unique_id_; } @@ -294,10 +284,6 @@ class HloModule { // sequential order of instructions for each non-fusion computation in the // module. absl::optional schedule_; - - // alias_config indicates the alias information of input/output buffers that - // are expected from the module. - HloInputOutputAliasConfig input_output_alias_config_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc index 2902a11a42..be3bee5975 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier.cc @@ -1220,8 +1220,6 @@ StatusOr HloVerifier::Run(HloModule* module) { TF_RETURN_IF_ERROR(module->schedule().Verify()); } - TF_RETURN_IF_ERROR(module->input_output_alias_config().Verify(*module)); - return false; } diff --git a/tensorflow/compiler/xla/shape_util.h b/tensorflow/compiler/xla/shape_util.h index 51cedce7f0..73f541d505 100644 --- a/tensorflow/compiler/xla/shape_util.h +++ b/tensorflow/compiler/xla/shape_util.h @@ -72,7 +72,7 @@ class ShapeIndex { void push_back(int64 value) { indices_.push_back(value); } void pop_back() { indices_.pop_back(); } - // push_front is O(n), but shapes don't usually have a ton of dimensions. + // push_front is O(n^2), but shapes don't usually have a ton of dimensions. void push_front(int64 value) { indices_.insert(indices_.begin(), value); } using container_type = absl::InlinedVector; -- cgit v1.2.3