1 | #include "taichi/ir/ir.h" |
---|---|
2 | #include "taichi/ir/transforms.h" |
3 | #include "taichi/ir/visitors.h" |
4 | |
5 | namespace taichi::lang { |
6 | |
7 | // This pass manipulates the id of statements so that they are successive values |
8 | // starting from 0 |
9 | class ReId : public BasicStmtVisitor { |
10 | public: |
11 | using BasicStmtVisitor::visit; |
12 | |
13 | int id_counter; |
14 | |
15 | ReId() : id_counter(0) { |
16 | allow_undefined_visitor = true; |
17 | invoke_default_visitor = true; |
18 | } |
19 | |
20 | void re_id(Stmt *stmt) { |
21 | stmt->id = id_counter++; |
22 | } |
23 | |
24 | void visit(Stmt *stmt) override { |
25 | re_id(stmt); |
26 | } |
27 | |
28 | void preprocess_container_stmt(Stmt *stmt) override { |
29 | re_id(stmt); |
30 | } |
31 | |
32 | static void run(IRNode *node) { |
33 | ReId instance; |
34 | node->accept(&instance); |
35 | } |
36 | }; |
37 | |
38 | namespace irpass { |
39 | void re_id(IRNode *root) { |
40 | ReId::run(root); |
41 | } |
42 | } // namespace irpass |
43 | |
44 | } // namespace taichi::lang |
45 |