Skip to content

Commit 116b001

Browse files
bowang007bowang007
committed
chore: apply linting
Signed-off-by: Bo Wang <[email protected]>
1 parent 3d39d7c commit 116b001

File tree

5 files changed

+90
-76
lines changed

5 files changed

+90
-76
lines changed

‎tests/core/partitioning/test_fallback_graph_output.cpp

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,17 @@
11
#include <string>
22
#include <unordered_set>
3-
#include "gtest/gtest.h"
4-
#include "torch/script.h"
53
#include "core/compiler.h"
4+
#include "gtest/gtest.h"
65
#include "tests/util/util.h"
7-
6+
#include "torch/script.h"
87

98
TEST(Partitioning, StitchSegmentedBlockCorrectly) {
109
torch::jit::script::Module mod;
1110
try {
12-
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
11+
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
1312
} catch (const c10::Error& e) {
14-
std::cerr << "error loading the model\n";
15-
return;
13+
std::cerr << "error loading the model\n";
14+
return;
1615
}
1716

1817
const std::vector<std::vector<int64_t>> input_shapes = {{3, 3, 16, 16}};
@@ -37,19 +36,19 @@ TEST(Partitioning, StitchSegmentedBlockCorrectly) {
3736
TEST(Partitioning, StitchSegmentedBlockCorrectlyEdge) {
3837
torch::jit::script::Module mod;
3938
try {
40-
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
39+
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
4140
} catch (const c10::Error& e) {
42-
std::cerr << "error loading the model\n";
43-
return;
41+
std::cerr << "error loading the model\n";
42+
return;
4443
}
4544

4645
const std::vector<std::vector<int64_t>> input_shapes = {{3, 3, 16, 16}};
4746
std::vector<torch::jit::IValue> jit_inputs_ivalues;
4847
std::vector<torch::jit::IValue> trt_inputs_ivalues;
4948
for (auto in_shape : input_shapes) {
50-
auto in = at::randint(5, in_shape, {at::kCUDA});
51-
jit_inputs_ivalues.push_back(in.clone());
52-
trt_inputs_ivalues.push_back(in.clone());
49+
auto in = at::randint(5, in_shape, {at::kCUDA});
50+
jit_inputs_ivalues.push_back(in.clone());
51+
trt_inputs_ivalues.push_back(in.clone());
5352
}
5453

5554
std::vector<trtorch::core::ir::InputRange> input_ranges{trtorch::core::ir::InputRange({3, 3, 16, 16})};
@@ -61,4 +60,3 @@ TEST(Partitioning, StitchSegmentedBlockCorrectlyEdge) {
6160
auto trt_results = trt_mod.forward(trt_inputs_ivalues).toTensor();
6261
ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results, trt_results, 2e-6));
6362
}
64-

‎tests/core/partitioning/test_segmentation.cpp

Lines changed: 43 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,27 @@
11
#include <string>
2+
#include "core/lowering/lowering.h"
3+
#include "core/partitioning/partitioning.h"
24
#include "gtest/gtest.h"
35
#include "tests/util/util.h"
46
#include "torch/script.h"
57
#include "trtorch/trtorch.h"
6-
#include "core/lowering/lowering.h"
7-
#include "core/partitioning/partitioning.h"
8-
98

10-
bool checkSegmentedBlockNumber(std::vector<trtorch::core::partitioning::SegmentedBlock>& segmented_blocks,
11-
trtorch::core::partitioning::SegmentedBlock::SegmentedBlockTarget target, int target_count) {
12-
for (auto &seg_block : segmented_blocks) {
9+
bool checkSegmentedBlockNumber(
10+
std::vector<trtorch::core::partitioning::SegmentedBlock>& segmented_blocks,
11+
trtorch::core::partitioning::SegmentedBlock::SegmentedBlockTarget target,
12+
int target_count) {
13+
for (auto& seg_block : segmented_blocks) {
1314
if (seg_block.target() == target) {
1415
target_count--;
1516
}
1617
}
1718
return target_count == 0;
1819
}
1920

20-
bool checkSegmentedBlockNodesMapping(std::vector<trtorch::core::partitioning::SegmentedBlock>& segmented_blocks,
21-
std::shared_ptr<torch::jit::Graph> g, std::vector<std::vector<int>> nodes_index) {
21+
bool checkSegmentedBlockNodesMapping(
22+
std::vector<trtorch::core::partitioning::SegmentedBlock>& segmented_blocks,
23+
std::shared_ptr<torch::jit::Graph> g,
24+
std::vector<std::vector<int>> nodes_index) {
2225
std::vector<torch::jit::Node*> graph_nodes;
2326
for (const auto n : g->nodes()) {
2427
if (n->kind() != torch::jit::prim::Constant) {
@@ -32,25 +35,27 @@ bool checkSegmentedBlockNodesMapping(std::vector<trtorch::core::partitioning::Se
3235
return false;
3336
}
3437
}
35-
if (seg_block_node_id != segmented_blocks[i].raw_nodes().size()) return false;
38+
if (seg_block_node_id != segmented_blocks[i].raw_nodes().size())
39+
return false;
3640
}
3741
return true;
3842
}
3943

4044
TEST(Partitioning, SegmentingGraphDefaultCorrectly) {
4145
torch::jit::script::Module mod;
4246
try {
43-
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
47+
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
4448
} catch (const c10::Error& e) {
45-
std::cerr << "error loading the model\n";
46-
return;
49+
std::cerr << "error loading the model\n";
50+
return;
4751
}
4852
auto graph_and_parameters = trtorch::core::lowering::Lower(mod, "forward");
4953
auto g = graph_and_parameters.first;
5054

5155
trtorch::core::partitioning::PartitionInfo partition_info;
5256
partition_info.enabled = true;
53-
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks = trtorch::core::partitioning::segment_graph(g, partition_info);
57+
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks =
58+
trtorch::core::partitioning::segment_graph(g, partition_info);
5459
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTensorRT, 2));
5560
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTorch, 1));
5661
ASSERT_TRUE(checkSegmentedBlockNodesMapping(segmented_blocks, g, {{0, 1, 2}, {3}, {4}}));
@@ -59,18 +64,19 @@ TEST(Partitioning, SegmentingGraphDefaultCorrectly) {
5964
TEST(Partitioning, SegmentingGraphWithMinBlockSizeCorrectly) {
6065
torch::jit::script::Module mod;
6166
try {
62-
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
67+
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
6368
} catch (const c10::Error& e) {
64-
std::cerr << "error loading the model\n";
65-
return;
69+
std::cerr << "error loading the model\n";
70+
return;
6671
}
6772
auto graph_and_parameters = trtorch::core::lowering::Lower(mod, "forward");
6873
auto g = graph_and_parameters.first;
6974

7075
trtorch::core::partitioning::PartitionInfo partition_info;
7176
partition_info.enabled = true;
7277
partition_info.min_block_size = 3;
73-
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks = trtorch::core::partitioning::segment_graph(g, partition_info);
78+
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks =
79+
trtorch::core::partitioning::segment_graph(g, partition_info);
7480
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTensorRT, 1));
7581
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTorch, 1));
7682
ASSERT_TRUE(checkSegmentedBlockNodesMapping(segmented_blocks, g, {{0, 1, 2}, {3, 4}}));
@@ -79,18 +85,19 @@ TEST(Partitioning, SegmentingGraphWithMinBlockSizeCorrectly) {
7985
TEST(Partitioning, SegmentingGraphWithForcedOPeCorrectly) {
8086
torch::jit::script::Module mod;
8187
try {
82-
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
88+
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
8389
} catch (const c10::Error& e) {
84-
std::cerr << "error loading the model\n";
85-
return;
90+
std::cerr << "error loading the model\n";
91+
return;
8692
}
8793
auto graph_and_parameters = trtorch::core::lowering::Lower(mod, "forward");
8894
auto g = graph_and_parameters.first;
8995

9096
trtorch::core::partitioning::PartitionInfo partition_info;
9197
partition_info.enabled = true;
9298
partition_info.forced_fallback_operators.push_back("aten::relu");
93-
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks = trtorch::core::partitioning::segment_graph(g, partition_info);
99+
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks =
100+
trtorch::core::partitioning::segment_graph(g, partition_info);
94101
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTensorRT, 3));
95102
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTorch, 2));
96103
ASSERT_TRUE(checkSegmentedBlockNodesMapping(segmented_blocks, g, {{0}, {1}, {2}, {3}, {4}}));
@@ -99,17 +106,18 @@ TEST(Partitioning, SegmentingGraphWithForcedOPeCorrectly) {
99106
TEST(Partitioning, SegmentingGraphDefaultCorrectlyEdge) {
100107
torch::jit::script::Module mod;
101108
try {
102-
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
109+
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
103110
} catch (const c10::Error& e) {
104-
std::cerr << "error loading the model\n";
105-
return;
111+
std::cerr << "error loading the model\n";
112+
return;
106113
}
107114
auto graph_and_parameters = trtorch::core::lowering::Lower(mod, "forward");
108115
auto g = graph_and_parameters.first;
109116

110117
trtorch::core::partitioning::PartitionInfo partition_info;
111118
partition_info.enabled = true;
112-
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks = trtorch::core::partitioning::segment_graph(g, partition_info);
119+
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks =
120+
trtorch::core::partitioning::segment_graph(g, partition_info);
113121
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTensorRT, 2));
114122
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTorch, 1));
115123
ASSERT_TRUE(checkSegmentedBlockNodesMapping(segmented_blocks, g, {{0, 1}, {2}, {3, 4, 5, 6}}));
@@ -118,18 +126,19 @@ TEST(Partitioning, SegmentingGraphDefaultCorrectlyEdge) {
118126
TEST(Partitioning, SegmentingGraphWithMinBlockSizeCorrectlyEdge) {
119127
torch::jit::script::Module mod;
120128
try {
121-
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
129+
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
122130
} catch (const c10::Error& e) {
123-
std::cerr << "error loading the model\n";
124-
return;
131+
std::cerr << "error loading the model\n";
132+
return;
125133
}
126134
auto graph_and_parameters = trtorch::core::lowering::Lower(mod, "forward");
127135
auto g = graph_and_parameters.first;
128136

129137
trtorch::core::partitioning::PartitionInfo partition_info;
130138
partition_info.enabled = true;
131139
partition_info.min_block_size = 3;
132-
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks = trtorch::core::partitioning::segment_graph(g, partition_info);
140+
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks =
141+
trtorch::core::partitioning::segment_graph(g, partition_info);
133142
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTensorRT, 1));
134143
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTorch, 1));
135144
ASSERT_TRUE(checkSegmentedBlockNodesMapping(segmented_blocks, g, {{0, 1, 2}, {3, 4, 5, 6}}));
@@ -138,18 +147,19 @@ TEST(Partitioning, SegmentingGraphWithMinBlockSizeCorrectlyEdge) {
138147
TEST(Partitioning, SegmentingGraphWithForcedOPeCorrectlyEdge) {
139148
torch::jit::script::Module mod;
140149
try {
141-
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
150+
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
142151
} catch (const c10::Error& e) {
143-
std::cerr << "error loading the model\n";
144-
return;
152+
std::cerr << "error loading the model\n";
153+
return;
145154
}
146155
auto graph_and_parameters = trtorch::core::lowering::Lower(mod, "forward");
147156
auto g = graph_and_parameters.first;
148157

149158
trtorch::core::partitioning::PartitionInfo partition_info;
150159
partition_info.enabled = true;
151160
partition_info.forced_fallback_operators.push_back("aten::relu");
152-
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks = trtorch::core::partitioning::segment_graph(g, partition_info);
161+
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks =
162+
trtorch::core::partitioning::segment_graph(g, partition_info);
153163
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTensorRT, 3));
154164
ASSERT_TRUE(checkSegmentedBlockNumber(segmented_blocks, trtorch::core::partitioning::SegmentedBlock::kTorch, 2));
155165
ASSERT_TRUE(checkSegmentedBlockNodesMapping(segmented_blocks, g, {{0, 1}, {2}, {3}, {4}, {5, 6}}));

‎tests/core/partitioning/test_shape_analysis.cpp

Lines changed: 20 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,19 @@
11
#include <string>
2-
#include "gtest/gtest.h"
3-
#include "torch/script.h"
42
#include "core/lowering/lowering.h"
53
#include "core/partitioning/partitioning.h"
64
#include "core/util/trt_util.h"
5+
#include "gtest/gtest.h"
6+
#include "torch/script.h"
77

8-
bool checkSegmentedBlockInputShape(std::vector<trtorch::core::partitioning::SegmentedBlock>& segmented_blocks, std::vector<std::vector<std::vector<int>>> in_shape) {
9-
if (segmented_blocks.size() != in_shape.size()) return false;
8+
bool checkSegmentedBlockInputShape(
9+
std::vector<trtorch::core::partitioning::SegmentedBlock>& segmented_blocks,
10+
std::vector<std::vector<std::vector<int>>> in_shape) {
11+
if (segmented_blocks.size() != in_shape.size())
12+
return false;
1013
for (size_t i = 0; i < segmented_blocks.size(); ++i) {
1114
auto cur_block_in_shapes = segmented_blocks[i].in_shape();
12-
if (cur_block_in_shapes.size() != in_shape[i].size()) return false;
15+
if (cur_block_in_shapes.size() != in_shape[i].size())
16+
return false;
1317
for (size_t j = 0; j < cur_block_in_shapes.size(); ++j) {
1418
auto cur_input_shape = trtorch::core::util::toVec(cur_block_in_shapes[j]);
1519
for (size_t k = 0; k < cur_input_shape.size(); ++k) {
@@ -36,17 +40,19 @@ TEST(Partitioning, InferSegmentedBlockShapeCorrectly) {
3640
partition_info.enabled = true;
3741
std::vector<trtorch::core::ir::InputRange> input_ranges{trtorch::core::ir::InputRange({3, 3, 16, 16})};
3842

39-
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks = trtorch::core::partitioning::Partition(g, input_ranges, partition_info);
40-
ASSERT_TRUE(checkSegmentedBlockInputShape(segmented_blocks, {{{3, 3, 16, 16}}, {{3, 16, 16, 16}}, {{3, 16, 16, 16}}}));
43+
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks =
44+
trtorch::core::partitioning::Partition(g, input_ranges, partition_info);
45+
ASSERT_TRUE(
46+
checkSegmentedBlockInputShape(segmented_blocks, {{{3, 3, 16, 16}}, {{3, 16, 16, 16}}, {{3, 16, 16, 16}}}));
4147
}
4248

4349
TEST(Partitioning, InferSegmentedBlockShapeCorrectlyEdge) {
4450
torch::jit::script::Module mod;
4551
try {
46-
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
52+
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
4753
} catch (const c10::Error& e) {
48-
std::cerr << "error loading the model\n";
49-
return;
54+
std::cerr << "error loading the model\n";
55+
return;
5056
}
5157
auto graph_and_parameters = trtorch::core::lowering::Lower(mod, "forward");
5258
auto g = graph_and_parameters.first;
@@ -55,6 +61,8 @@ TEST(Partitioning, InferSegmentedBlockShapeCorrectlyEdge) {
5561
partition_info.enabled = true;
5662
std::vector<trtorch::core::ir::InputRange> input_ranges{trtorch::core::ir::InputRange({3, 3, 16, 16})};
5763

58-
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks = trtorch::core::partitioning::Partition(g, input_ranges, partition_info);
59-
ASSERT_TRUE(checkSegmentedBlockInputShape(segmented_blocks, {{{3, 3, 16, 16}}, {{3, 32, 16, 16}}, {{3, 32, 16, 16}, {3, 16, 16, 16}}}));
64+
std::vector<trtorch::core::partitioning::SegmentedBlock> segmented_blocks =
65+
trtorch::core::partitioning::Partition(g, input_ranges, partition_info);
66+
ASSERT_TRUE(checkSegmentedBlockInputShape(
67+
segmented_blocks, {{{3, 3, 16, 16}}, {{3, 32, 16, 16}}, {{3, 32, 16, 16}, {3, 16, 16, 16}}}));
6068
}

‎tests/core/partitioning/test_stitched_graph.cpp

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
#include <string>
22
#include <unordered_set>
3-
#include "gtest/gtest.h"
4-
#include "torch/script.h"
53
#include "core/compiler.h"
64
#include "core/util/trt_util.h"
5+
#include "gtest/gtest.h"
6+
#include "torch/script.h"
77

88
bool checkAllInputsExistInStitchedGraph(std::shared_ptr<torch::jit::Graph> g) {
99
std::unordered_set<torch::jit::Value*> available_values;
@@ -25,10 +25,10 @@ bool checkAllInputsExistInStitchedGraph(std::shared_ptr<torch::jit::Graph> g) {
2525
TEST(Partitioning, StitchSegmentedBlockCorrectly) {
2626
torch::jit::script::Module mod;
2727
try {
28-
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
28+
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
2929
} catch (const c10::Error& e) {
30-
std::cerr << "error loading the model\n";
31-
return;
30+
std::cerr << "error loading the model\n";
31+
return;
3232
}
3333

3434
std::vector<trtorch::core::ir::InputRange> input_ranges{trtorch::core::ir::InputRange({3, 3, 16, 16})};
@@ -39,14 +39,13 @@ TEST(Partitioning, StitchSegmentedBlockCorrectly) {
3939
ASSERT_TRUE(checkAllInputsExistInStitchedGraph(g));
4040
}
4141

42-
4342
TEST(Partitioning, StitchSegmentedBlockCorrectlyEdge) {
4443
torch::jit::script::Module mod;
4544
try {
46-
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
45+
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
4746
} catch (const c10::Error& e) {
48-
std::cerr << "error loading the model\n";
49-
return;
47+
std::cerr << "error loading the model\n";
48+
return;
5049
}
5150

5251
std::vector<trtorch::core::ir::InputRange> input_ranges{trtorch::core::ir::InputRange({3, 3, 16, 16})};

‎tests/core/partitioning/test_tensorrt_conversion.cpp

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
#include <string>
2-
#include "gtest/gtest.h"
3-
#include "torch/script.h"
42
#include "core/compiler.h"
53
#include "core/util/trt_util.h"
4+
#include "gtest/gtest.h"
5+
#include "torch/script.h"
66

77
int count_trt_engines(std::shared_ptr<torch::jit::Graph> g) {
88
int count = 0;
@@ -17,10 +17,10 @@ int count_trt_engines(std::shared_ptr<torch::jit::Graph> g) {
1717
TEST(Partitioning, ConvertSegmentedBlockCorrectly) {
1818
torch::jit::script::Module mod;
1919
try {
20-
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
20+
mod = torch::jit::load("tests/core/partitioning/test_base_model.jit");
2121
} catch (const c10::Error& e) {
22-
std::cerr << "error loading the model\n";
23-
return;
22+
std::cerr << "error loading the model\n";
23+
return;
2424
}
2525

2626
std::vector<trtorch::core::ir::InputRange> input_ranges{trtorch::core::ir::InputRange({3, 3, 16, 16})};
@@ -32,14 +32,13 @@ TEST(Partitioning, ConvertSegmentedBlockCorrectly) {
3232
ASSERT_TRUE(count == 2);
3333
}
3434

35-
3635
TEST(Partitioning, ConvertSegmentedBlockCorrectlyEdge) {
3736
torch::jit::script::Module mod;
3837
try {
39-
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
38+
mod = torch::jit::load("tests/core/partitioning/test_edge_model.jit");
4039
} catch (const c10::Error& e) {
41-
std::cerr << "error loading the model\n";
42-
return;
40+
std::cerr << "error loading the model\n";
41+
return;
4342
}
4443

4544
std::vector<trtorch::core::ir::InputRange> input_ranges{trtorch::core::ir::InputRange({3, 3, 16, 16})};

0 commit comments

Comments
 (0)