forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathprofiling_record.cpp
91 lines (77 loc) · 2.44 KB
/
profiling_record.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
#include <torch/csrc/jit/profiling_record.h>
namespace torch {
namespace jit {
ProfilingRecord::ProfilingRecord(std::shared_ptr<Graph> g)
: profiled_graph_(std::move(g)), profiling_count_(3) {}
ProfileOp* ProfilingRecord::createProfileNode(
const std::function<void(Stack&)>& fp,
at::ArrayRef<Value*> inputs) {
auto pn = new ProfileOp(profiled_graph_.get(), fp);
for (auto in : inputs) {
pn->addInput(in);
}
return pn;
}
void ProfilingRecord::instrumentBlock(Block* block) {
for (auto it = block->nodes().begin(); it != block->nodes().end(); ++it) {
auto n = *it;
for (auto i : n->inputs()) {
if (!i->type()->isSubclass(TypeKind::TensorType) ||
i->node()->kind() == prim::profile) {
continue;
}
auto pn = createProfileNode(nullptr, {i});
auto pno = pn->addOutput();
pno->setType(i->type());
std::function<void(Stack&)> shape_profiler = [this, pno](Stack& stack) {
IValue t;
pop(stack, t);
if (t.isTensor()) {
auto pttp = ProfiledTensorType::create(t.toTensor());
std::lock_guard<std::mutex> lock(this->mutex_);
if (pno->type()->isSubclass(TypeKind::ProfiledTensorType)) {
auto type = pno->type()->cast<ProfiledTensorType>();
pno->setType(type->merge(pttp));
} else {
pno->setType(pttp);
}
}
// passing t through
push(stack, t);
};
pn->setCallback(shape_profiler);
pn->insertBefore(n);
n->replaceInputWith(i, pn->output());
}
for (auto b : n->blocks()) {
instrumentBlock(b);
}
}
}
std::unique_ptr<ProfilingRecord> ProfilingRecord::instrumentGraph(
const std::shared_ptr<Graph>& graph) {
auto new_g = graph->copy();
auto pr = std::unique_ptr<ProfilingRecord>(new ProfilingRecord(new_g));
auto raw_pr = pr.get();
pr->instrumentBlock(new_g->block());
std::function<void(Stack&)> counter = [raw_pr](Stack&) {
std::lock_guard<std::mutex> lock(raw_pr->mutex_);
if (raw_pr->profiling_count_ > 0)
{
raw_pr->profiling_count_--;
}
};
auto pop = pr->createProfileNode(counter, {});
new_g->appendNode(pop);
return pr;
}
ProfiledTensorTypePtr ProfilingRecord::toProfiledTensorTypePtr(
const IValue& ival) {
if (ival.isTensor()) {
auto tensor = ival.toTensor();
return ProfiledTensorType::create(tensor);
}
return {nullptr};
}
} // namespace jit
} // namespace torch