1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" 17 18 #include <map> 19 #include <memory> 20 #include <string> 21 22 #include "tensorflow/compiler/xla/literal_util.h" 23 #include "tensorflow/compiler/xla/service/hlo_casting_utils.h" 24 #include "tensorflow/compiler/xla/service/hlo_creation_utils.h" 25 #include "tensorflow/compiler/xla/service/hlo_instructions.h" 26 #include "tensorflow/compiler/xla/service/hlo_matchers.h" 27 #include "tensorflow/compiler/xla/service/hlo_opcode.h" 28 #include "tensorflow/compiler/xla/service/instruction_fusion.h" 29 #include "tensorflow/compiler/xla/shape_util.h" 30 #include "tensorflow/compiler/xla/test.h" 31 #include "tensorflow/compiler/xla/test_helpers.h" 32 #include "tensorflow/compiler/xla/tests/hlo_test_base.h" 33 #include "tensorflow/compiler/xla/xla_data.pb.h" 34 #include "tensorflow/core/platform/logging.h" 35 #include "tensorflow/core/platform/test.h" 36 37 namespace op = xla::testing::opcode_matchers; 38 39 namespace xla { 40 namespace { 41 42 using ::testing::UnorderedElementsAre; 43 using ::testing::UnorderedElementsAreArray; 44 45 class TuplePointsToAnalysisTest : public HloTestBase { 46 protected: 47 // Builds a module with the given entry computation and runs points to 48 // analysis. BuildModuleAndRunAnalysis(std::unique_ptr<HloComputation> computation)49 void BuildModuleAndRunAnalysis(std::unique_ptr<HloComputation> computation) { 50 BuildModule(std::move(computation)); 51 RunAnalysis(); 52 } 53 BuildModule(std::unique_ptr<HloComputation> computation)54 void BuildModule(std::unique_ptr<HloComputation> computation) { 55 module_ = CreateNewVerifiedModule(); 56 module_->AddEntryComputation(std::move(computation)); 57 } 58 RunAnalysis()59 void RunAnalysis() { 60 CHECK_NOTNULL(module_.get()); 61 points_to_analysis_ = TuplePointsToAnalysis::Run(module_.get()).value(); 62 } 63 64 // Returns the LogicalBuffer defined at the given instruction and 65 // index. CHECKs if no buffer is defined at that point. GetBuffer(const HloInstruction * instruction,const ShapeIndex & index)66 const LogicalBuffer* const GetBuffer(const HloInstruction* instruction, 67 const ShapeIndex& index) { 68 const auto& pointed_to = 69 points_to_analysis_->GetPointsToSet(instruction).element(index); 70 CHECK_EQ(1, pointed_to.size()); 71 CHECK_EQ(instruction, pointed_to[0]->instruction()); 72 CHECK(index == pointed_to[0]->index()); 73 return pointed_to[0]; 74 } 75 76 // Checks that the given points-to set contains exactly (unordered) the given 77 // LogicalBuffers. ExpectHasBuffers(const PointsToSet::BufferList & points_to_set,absl::Span<const LogicalBuffer * const> buffers)78 void ExpectHasBuffers(const PointsToSet::BufferList& points_to_set, 79 absl::Span<const LogicalBuffer* const> buffers) { 80 std::vector<const LogicalBuffer*> vec(buffers.begin(), buffers.end()); 81 EXPECT_THAT(points_to_set, UnorderedElementsAreArray(vec)); 82 } 83 84 // Checks that the given points-to set contains exactly (unordered) the 85 // top-level buffers of the given instructions. ExpectHasTopLevelBuffers(const PointsToSet::BufferList & points_to_set,absl::Span<HloInstruction * const> instructions)86 void ExpectHasTopLevelBuffers( 87 const PointsToSet::BufferList& points_to_set, 88 absl::Span<HloInstruction* const> instructions) { 89 PointsToSet::BufferList buffers; 90 for (auto instruction : instructions) { 91 buffers.push_back(GetBuffer(instruction, /*index=*/{})); 92 } 93 ExpectHasBuffers(points_to_set, buffers); 94 } 95 96 // Overload which takes a set instead of a vector. ExpectHasTopLevelBuffers(const PointsToSet::BufferSet & points_to_set,absl::Span<HloInstruction * const> instructions)97 void ExpectHasTopLevelBuffers( 98 const PointsToSet::BufferSet& points_to_set, 99 absl::Span<HloInstruction* const> instructions) { 100 ExpectHasTopLevelBuffers( 101 PointsToSet::BufferList(points_to_set.begin(), points_to_set.end()), 102 instructions); 103 } 104 105 // Checks that the buffer defined at the given instruction and index has 106 // aliases which are exactly (unordered) the given instruction/index pairs. ExpectHasBufferAliases(const HloInstruction * instruction,const ShapeIndex & index,absl::Span<const std::pair<HloInstruction *,ShapeIndex>> expected)107 void ExpectHasBufferAliases( 108 const HloInstruction* instruction, const ShapeIndex& index, 109 absl::Span<const std::pair<HloInstruction*, ShapeIndex>> expected) { 110 const LogicalBuffer* buffer = 111 points_to_analysis_->GetBufferDefinedAt(instruction, index) 112 .ValueOrDie(); 113 std::vector<BufferAlias> expected_aliases; 114 expected_aliases.reserve(expected.size()); 115 for (auto& pair : expected) { 116 expected_aliases.push_back(BufferAlias(pair.first, pair.second)); 117 } 118 EXPECT_THAT(points_to_analysis_->GetBufferAliases(*buffer), 119 UnorderedElementsAreArray(expected_aliases)); 120 } 121 122 std::unique_ptr<HloModule> module_; 123 std::unique_ptr<TuplePointsToAnalysis> points_to_analysis_; 124 }; 125 126 } // namespace 127 } // namespace xla 128