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/hlo_execution_profile.h" 17 #include "tensorflow/compiler/xla/service/hlo_cost_analysis.h" 18 #include "tensorflow/compiler/xla/tests/hlo_test_base.h" 19 20 namespace xla { 21 namespace { 22 23 class HloExecutionProfileTest : public HloTestBase { 24 protected: 25 static constexpr int64 kInstructionCyclesIndex = 0; 26 static constexpr int64 kInstructionNameIndex = 19; 27 }; 28 29 // Splits `lines` into a sequence of lines delimited by newlines and then split 30 // each of those lines into a sequence of words delimited by spaces. Filter out 31 // empty words. 32 std::vector<std::vector<string>> SplitIntoLinesAndWords( 33 tensorflow::StringPiece lines) { 34 std::vector<std::vector<string>> result; 35 for (const string& line : tensorflow::str_util::Split(lines, '\n')) { 36 std::vector<string> words; 37 for (const string& word : tensorflow::str_util::Split(line, ' ')) { 38 if (!word.empty()) { 39 words.push_back(word); 40 } 41 } 42 result.push_back(std::move(words)); 43 } 44 45 return result; 46 } 47 48 TEST_F(HloExecutionProfileTest, Basic) { 49 std::unique_ptr<HloModule> hlo_module = CreateNewModule(); 50 51 HloComputation::Builder builder(TestName()); 52 Shape shape = ShapeUtil::MakeShape(F32, {30, 30}); 53 HloInstruction* param_lhs = 54 builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "lhs")); 55 HloInstruction* param_rhs = 56 builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "rhs")); 57 HloInstruction* add_instruction = 58 builder.AddInstruction(HloInstruction::CreateBinary( 59 shape, HloOpcode::kAdd, param_lhs, param_rhs)); 60 HloInstruction* dot_instruction = 61 builder.AddInstruction(HloInstruction::CreateBinary( 62 shape, HloOpcode::kDot, param_lhs, add_instruction)); 63 64 hlo_module->AddEntryComputation(builder.Build()); 65 66 auto shape_size_function = [&](const Shape& shape) { 67 const int64 pointer_size = 8; 68 if (ShapeUtil::IsOpaque(shape)) { 69 return pointer_size; 70 } 71 return ShapeUtil::ByteSizeOf(shape, pointer_size); 72 }; 73 74 HloCostAnalysis cost_analysis(shape_size_function); 75 HloProfileIndexMap profile_index_map(*hlo_module); 76 std::unique_ptr<HloProfilePrinterData> profile_printer = 77 CreateHloProfilePrinterData(profile_index_map, cost_analysis); 78 HloExecutionProfile execution_profile(profile_printer.get(), 79 &profile_index_map); 80 81 const int64 add_cycles = 1000; 82 const int64 dot_cycles = 4000; 83 84 execution_profile.SetCyclesTakenBy(add_instruction, add_cycles); 85 execution_profile.SetCyclesTakenBy(dot_instruction, dot_cycles); 86 87 string rendered_profile = execution_profile.ToString( 88 backend().default_stream_executor()->GetDeviceDescription()); 89 std::vector<std::vector<string>> lines_and_words = 90 SplitIntoLinesAndWords(rendered_profile); 91 ASSERT_EQ(lines_and_words.size(), 8); 92 93 const std::vector<string>& line_2 = lines_and_words[2]; 94 const std::vector<string>& line_3 = lines_and_words[3]; 95 96 EXPECT_EQ(line_2[kInstructionCyclesIndex], std::to_string(dot_cycles)); 97 EXPECT_EQ(line_2[kInstructionNameIndex], '%' + dot_instruction->name()); 98 99 EXPECT_EQ(line_3[kInstructionCyclesIndex], std::to_string(add_cycles)); 100 EXPECT_EQ(line_3[kInstructionNameIndex], '%' + add_instruction->name()); 101 } 102 } // namespace 103 } // namespace xla 104
