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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
|
#include <gtest/gtest.h>
#include <test/cpp/jit/test_utils.h>
#include <torch/csrc/jit/ir/ir.h>
#include <torch/csrc/jit/ir/irparser.h>
#include <torch/csrc/jit/passes/peephole.h>
namespace torch {
namespace jit {
TEST(PeepholeOptimizeTest, IsAndIsNot)
// test is / is not none optimization
{
auto graph = std::make_shared<Graph>();
parseIR(
R"IR(
graph(%0 : int):
%1 : None = prim::Constant()
%2 : bool = aten::__is__(%0, %1)
%3 : bool = aten::__isnot__(%0, %1)
return (%2, %3)
)IR",
graph.get());
PeepholeOptimize(graph);
testing::FileCheck()
.check_not("aten::__is__")
->check_not("aten::__isnot__")
->run(*graph);
}
TEST(PeepholeOptimizeTest, IsAndIsNot2) {
auto graph = std::make_shared<Graph>();
parseIR(
R"IR(
graph(%0: int?):
%1 : None = prim::Constant()
%2 : bool = aten::__is__(%0, %1)
%3 : bool = aten::__isnot__(%0, %1)
return (%2, %3)
)IR",
graph.get());
PeepholeOptimize(graph);
testing::FileCheck()
.check("aten::__is__")
->check("aten::__isnot__")
->run(*graph);
}
TEST(PeepholeOptimizeTest, IsAndIsNot3) {
auto graph = std::make_shared<Graph>();
parseIR(
R"IR(
graph(%0: int?):
%1 : Tensor = prim::AutogradZero()
%2 : None = prim::Constant()
%4 : bool = aten::__is__(%0, %1)
%5 : bool = aten::__isnot__(%1, %2)
return (%4, %5)
)IR",
graph.get());
PeepholeOptimize(graph);
testing::FileCheck()
.check("aten::__is__")
->check_not("aten::__isnot__")
->run(*graph);
}
TEST(PeepholeOptimizeTest, UnwrapOptional)
// test unwrap optional
{
auto graph = std::make_shared<Graph>();
parseIR(
R"IR(
graph():
%1 : Float(*, *, *) = prim::Constant()
%2 : bool = aten::_unwrap_optional(%1)
%3 : bool = prim::unchecked_unwrap_optional(%1)
return (%2, %3)
)IR",
graph.get());
PeepholeOptimize(graph);
testing::FileCheck().check_not("unwrap")->run(*graph);
}
TEST(PeepholeOptimizeTest, UnwrapOptional2) {
auto graph = std::make_shared<Graph>();
parseIR(
R"IR(
graph(%1 : Float(*, *, *)?):
%2 : bool = aten::_unwrap_optional(%1)
%3 : bool = prim::unchecked_unwrap_optional(%1)
return (%2, %3)
)IR",
graph.get());
PeepholeOptimize(graph);
testing::FileCheck().check_count("unwrap", 2)->run(*graph);
}
TEST(PeepholeOptimizeTest, AddMMFusion) {
auto graph = std::make_shared<Graph>();
parseIR(
R"IR(
graph(
%0 : Float(2, 3, 4),
%1 : Float(2, 3, 4),
%2 : Float(1, 1, 1)):
%3 : int = prim::Constant[value=1]()
%4 : Tensor = aten::mm(%0, %1)
%5 : Tensor = aten::add(%4, %2, %3)
%6 : Tensor = aten::add(%5, %2, %3)
return (%6)
)IR",
graph.get());
FuseAddMM(graph);
testing::FileCheck().check("addmm")->run(*graph);
}
} // namespace jit
} // namespace torch
|
