1#include <torch/csrc/jit/passes/annotate_warns.h>
2
3#include <atomic>
4
5namespace torch {
6namespace jit {
7
8void AnnotateWarns(Block* b) {
9 static std::atomic<int64_t> idx(0);
10 for (Node* n : b->nodes()) {
11 for (Block* child_b : n->blocks()) {
12 AnnotateWarns(child_b);
13 }
14
15 if (n->kind() != aten::warn) {
16 continue;
17 }
18
19 n->i_(attr::warn_id, idx);
20 idx++;
21 }
22}
23
24void AnnotateWarns(const std::shared_ptr<Graph>& graph) {
25 AnnotateWarns(graph->block());
26}
27
28} // namespace jit
29} // namespace torch
30