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
76
77
78
79
80
81
82
83
84
85
|
/* 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/gpu/layout_assignment.h"
#include "tensorflow/compiler/xla/layout_util.h"
#include "tensorflow/compiler/xla/service/computation_layout.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
#include "tensorflow/compiler/xla/shape_layout.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
namespace xla {
namespace gpu {
namespace {
using LayoutAssignmentTest = HloTestBase;
TEST_F(LayoutAssignmentTest, Elementwise) {
Shape ashape = ShapeUtil::MakeShape(F32, {42, 12});
Shape ashape_in_row_major(ashape);
Shape ashape_in_col_major(ashape);
*ashape_in_row_major.mutable_layout() = LayoutUtil::MakeLayout({1, 0});
*ashape_in_col_major.mutable_layout() = LayoutUtil::MakeLayout({0, 1});
// Enumerate all possible combinations of layouts.
for (const Shape& lhs_shape_with_layout :
{ashape_in_row_major, ashape_in_col_major}) {
for (const Shape& rhs_shape_with_layout :
{ashape_in_row_major, ashape_in_col_major}) {
for (const Shape& result_shape_with_layout :
{ashape_in_row_major, ashape_in_col_major}) {
// GpuLayoutAssignment should assign the same layout to "add" and its
// two operands.
auto builder = HloComputation::Builder(TestName());
auto x = builder.AddInstruction(
HloInstruction::CreateParameter(0, ashape, "x"));
auto y = builder.AddInstruction(
HloInstruction::CreateParameter(1, ashape, "y"));
auto add = builder.AddInstruction(
HloInstruction::CreateBinary(ashape, HloOpcode::kAdd, x, y));
HloModule module(TestName());
HloComputation* computation =
module.AddEntryComputation(builder.Build(add));
ComputationLayout computation_layout(
computation->ComputeProgramShape());
*computation_layout.mutable_parameter_layout(0) =
ShapeLayout(lhs_shape_with_layout);
*computation_layout.mutable_parameter_layout(1) =
ShapeLayout(rhs_shape_with_layout);
*computation_layout.mutable_result_layout() =
ShapeLayout(result_shape_with_layout);
GpuLayoutAssignment layout_assignment(&computation_layout);
EXPECT_TRUE(layout_assignment.Run(&module).ValueOrDie());
for (const HloInstruction* operand : add->operands()) {
EXPECT_TRUE(LayoutUtil::Equal(add->shape().layout(),
operand->shape().layout()));
}
}
}
}
}
} // namespace
} // namespace gpu
} // namespace xla
|
