1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
|
/* 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_execution_profile.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/xla/service/hlo_cost_analysis.h"
#include "tensorflow/compiler/xla/service/hlo_parser.h"
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
namespace xla {
namespace {
using absl::StrCat;
using ::testing::AllOf;
using ::testing::ContainsRegex;
class HloExecutionProfileTest : public HloTestBase {};
TEST_F(HloExecutionProfileTest, Basic) {
auto hlo_module = ParseHloString(R"(
HloModule test_module
ENTRY entry_computation {
lhs = f32[30,30]{1,0} parameter(0)
rhs = f32[30,30]{1,0} parameter(1)
add = f32[30,30]{1,0} add(lhs, rhs)
ROOT dot = f32[30,30]{1,0} dot(lhs, add), lhs_contracting_dims={1}, rhs_contracting_dims={0}
})")
.ValueOrDie();
const HloInstruction* dot_instruction =
hlo_module->entry_computation()->root_instruction();
const HloInstruction* add_instruction = dot_instruction->operand(1);
Shape shape = ShapeUtil::MakeShape(F32, {30, 30});
auto shape_size_function = [&](const Shape& shape) {
const int64 pointer_size = 8;
if (ShapeUtil::IsOpaque(shape)) {
return pointer_size;
}
return ShapeUtil::ByteSizeOf(shape, pointer_size);
};
HloCostAnalysis cost_analysis(shape_size_function);
HloProfileIndexMap profile_index_map(*hlo_module);
std::unique_ptr<HloProfilePrinterData> profile_printer =
CreateHloProfilePrinterData(profile_index_map, cost_analysis);
HloExecutionProfile execution_profile(profile_printer.get(),
&profile_index_map);
const int64 add_cycles = 1000;
const int64 dot_cycles = 4000;
execution_profile.SetCyclesTakenBy(add_instruction, add_cycles);
execution_profile.SetCyclesTakenBy(dot_instruction, dot_cycles);
EXPECT_THAT(execution_profile.ToString(
backend().default_stream_executor()->GetDeviceDescription()),
AllOf(ContainsRegex(StrCat(dot_cycles, R"(\b.*%)",
dot_instruction->name())),
ContainsRegex(StrCat(add_cycles, R"(\b.*%)",
add_instruction->name()))));
}
} // namespace
} // namespace xla
|
