-
Notifications
You must be signed in to change notification settings - Fork 344
/
fuse_addmm_branches.cpp
96 lines (79 loc) · 3 KB
/
fuse_addmm_branches.cpp
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
#include "torch/csrc/jit/ir/alias_analysis.h"
#include "torch/csrc/jit/jit_log.h"
#include "torch/csrc/jit/passes/constant_propagation.h"
#include "torch/csrc/jit/passes/dead_code_elimination.h"
#include "torch/csrc/jit/passes/guard_elimination.h"
#include "torch/csrc/jit/passes/peephole.h"
#include "torch/csrc/jit/runtime/graph_executor.h"
#include "core/util/prelude.h"
#include <vector>
namespace torch_tensorrt {
namespace core {
namespace lowering {
namespace passes {
namespace {
using namespace torch::jit;
struct AddMMBranchFusion {
AddMMBranchFusion(std::shared_ptr<Graph> graph) : graph_(std::move(graph)) {}
void run() {
findAddMMVariantsNodes(graph_->block());
torch::jit::EliminateDeadCode(graph_);
LOG_GRAPH("Post aten::addmm branch fusion: " << *graph_);
}
private:
bool isAddMMVariantsNode(Node* n) {
/// Check if this Node hosts a pattern like so:
/// %ret : Tensor = prim::If(%622)
/// block0():
/// %ret.1 : Tensor = aten::addmm(%self.fc.bias, %x9.1, %3677, %3, %3)
/// -> (%ret.1)
/// block1():
/// %output.1 : Tensor = aten::matmul(%x9.1, %3677)
/// %output0.1 : Tensor = aten::add(%output.1, %self.fc.bias, %3)
/// -> (%output0.1)
if (n->blocks().size() != 2) {
return false;
}
auto arm1 = n->blocks()[0];
auto arm2 = n->blocks()[1];
auto arm1_start = arm1->nodes().begin();
auto arm2_start = arm2->nodes().begin();
if ((*arm1_start)->kind().toQualString() == std::string("aten::addmm") &&
(*(++arm1_start))->kind() == prim::Return &&
(*arm2_start)->kind().toQualString() == std::string("aten::matmul") &&
(*(++arm2_start))->kind().toQualString() == std::string("aten::add") &&
(*(++arm2_start))->kind() == prim::Return) {
// Make sure that block0 is solely just the aten::addmm op and block1 is matmul + add
return true;
}
return false;
}
void findAddMMVariantsNodes(Block* b) {
for (auto it = b->nodes().begin(); it != b->nodes().end(); it++) {
auto n = *it;
if (n->kind() == prim::If && isAddMMVariantsNode(n)) {
LOG_GRAPH("Found that node " << *n << " is an AddMM variants node (FuseAddMMBranches)" << std::endl);
auto arm1 = n->blocks()[0];
auto arm1_start = arm1->nodes().begin();
auto input_values = (*arm1_start)->inputs();
auto new_addmm_node = b->owningGraph()->create(c10::Symbol::fromQualString("aten::addmm"), input_values, 1);
n->replaceAllUsesWith(new_addmm_node);
auto old_insert_point = b->owningGraph()->insertPoint();
b->owningGraph()->setInsertPoint(n);
b->owningGraph()->insertNode(new_addmm_node);
b->owningGraph()->setInsertPoint(old_insert_point);
it.destroyCurrent();
}
}
}
std::shared_ptr<Graph> graph_;
};
} // namespace
void FuseAddMMBranches(std::shared_ptr<Graph> graph) {
AddMMBranchFusion ammbf(std::move(graph));
ammbf.run();
}
} // namespace passes
} // namespace lowering
} // namespace core
} // namespace torch_tensorrt