1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7    http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
20namespace xla {
21namespace {
22
23class 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.
32std::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
48TEST_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